{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JndnmDMp66FL"
      },
      "source": [
        "Copyright 2018 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hMqWDc_m6rUC"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1KzHz-W-T504"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "import jax\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt\n",
        "import plotly.graph_objects as go\n",
        "import functools\n",
        "import jax.experimental.optimizers\n",
        "import time\n",
        "import flax\n",
        "import flax.linen as nn\n",
        "from typing import Sequence, Callable\n",
        "from IPython.display import clear_output\n",
        "import cv2\n",
        "import imageio\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qK31kBAtIJ3n"
      },
      "outputs": [],
      "source": [
        "def linear_to_srgb(linear):\n",
        "  \"\"\"Assumes `linear` is in [0, 1], see https://en.wikipedia.org/wiki/SRGB.\"\"\"\n",
        "  eps = jnp.finfo(jnp.float32).eps\n",
        "  srgb0 = 323 / 25 * linear\n",
        "  srgb1 = (211 * jnp.maximum(eps, linear)**(5 / 12) - 11) / 200\n",
        "  return jnp.where(linear \u003c= 0.0031308, srgb0, srgb1)\n",
        "\n",
        "def read_envmap(filename):\n",
        "  with open(filename, 'rb') as f:\n",
        "    return imageio.imread(f, 'exr')\n",
        "\n",
        "#envmap_linear = read_envmap(f'{DIRECTORY}/ninomaru_teien_4k.exr')\n",
        "envmap_linear = read_envmap(f'{DIRECTORY}/spruit_sunrise_4k.exr') / 3.0\n",
        "envmap_srgb = linear_to_srgb(envmap_linear)\n",
        "\n",
        "plt.imshow(envmap_srgb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hvG8F8ADKMHb"
      },
      "outputs": [],
      "source": [
        "envmap_H = 50\n",
        "envmap_W = envmap_H * 2 - 1\n",
        "\n",
        "#envmap_gt = linear_to_srgb(cv2.resize(envmap_linear, dsize=(envmap_W, envmap_H), interpolation=cv2.INTER_AREA))\n",
        "envmap_gt = cv2.resize(envmap_linear, dsize=(envmap_W, envmap_H), interpolation=cv2.INTER_AREA)\n",
        "plt.imshow(envmap_gt)\n",
        "plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hrCaejuLFcpI"
      },
      "outputs": [],
      "source": [
        "omega_phi, omega_theta = jnp.meshgrid(jnp.linspace(0.0, 2*jnp.pi, envmap_W+1)[:-1],\n",
        "                                      jnp.linspace(0.0,   jnp.pi, envmap_H+1)[:-1])\n",
        "\n",
        "dtheta_dphi = (omega_theta[1, 1] - omega_theta[0, 0]) * (omega_phi[1, 1] - omega_phi[0, 0])\n",
        "\n",
        "omega_theta = omega_theta.flatten()\n",
        "omega_phi = omega_phi.flatten()\n",
        "\n",
        "omega_x = jnp.sin(omega_theta) * jnp.cos(omega_phi)\n",
        "omega_y = jnp.sin(omega_theta) * jnp.sin(omega_phi)\n",
        "omega_z = jnp.cos(omega_theta)\n",
        "omega_xyz = jnp.stack([omega_x,\n",
        "                       omega_y,\n",
        "                       omega_z], axis=-1)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gJbvom3l620T"
      },
      "outputs": [],
      "source": [
        "def mse_to_psnr(mse):\n",
        "  \"\"\"Compute PSNR given an MSE (we assume the maximum pixel value is 1).\"\"\"\n",
        "  return -10. / jnp.log(10.) * jnp.log(mse)\n",
        "  \n",
        "def get_rays(H, W, focal, c2w, rand_ort=False, key=None):\n",
        "  \"\"\"\n",
        "  c2w: 4x4 matrix\n",
        "  output: two arrays of shape [H, W, 3]\n",
        "  \"\"\"\n",
        "  j, i = jnp.meshgrid(jnp.arange(W, dtype=jnp.float32),\n",
        "                      jnp.arange(H, dtype=jnp.float32))\n",
        "  \n",
        "  if rand_ort:\n",
        "    k1, k2 = random.split(key)\n",
        "\n",
        "    i += jax.random.uniform(k1, shape=(H, W)) - 0.5\n",
        "    j += jax.random.uniform(k2, shape=(H, W)) - 0.5\n",
        "      \n",
        "  dirs = jnp.stack([ (j.flatten()-0.5*W)/focal,\n",
        "                    -(i.flatten()-0.5*H)/focal,\n",
        "                    -jnp.ones((H*W,), dtype=jnp.float32)], -1)  # shape [HW, 3]\n",
        "  \n",
        "  rays_d = dirs @ c2w[:3, :3].T  # shape [HW, 3]\n",
        "  rays_o = c2w[:3,-1:].T.repeat(H*W, 0)\n",
        "  return rays_o.reshape(H, W, 3), rays_d.reshape(H, W, 3)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bbHRuTQL7GZS"
      },
      "outputs": [],
      "source": [
        "def parse_bin(s):\n",
        "  return int(s[1:], 2) / 2.**(len(s) - 1)\n",
        "\n",
        "\n",
        "def phi2(i):\n",
        "  return parse_bin('.' + f'{i:b}'[::-1])\n",
        "\n",
        "def nice_uniform(N):\n",
        "  u = []\n",
        "  v = []\n",
        "  for i in range(N):\n",
        "    u.append(i / float(N))\n",
        "    v.append(phi2(i))\n",
        "    #pts.append((i/float(N), phi2(i)))\n",
        "\n",
        "  return u, v\n",
        "\n",
        "def nice_uniform_spherical(N, hemisphere=True):\n",
        "  \"\"\"implementation of http://holger.dammertz.org/stuff/notes_HammersleyOnHemisphere.html\"\"\"\n",
        "  u, v = nice_uniform(N)\n",
        "\n",
        "  theta = np.arccos(1.0 - np.array(u)) * (2.0 - int(hemisphere))\n",
        "  phi   = 2.0 * np.pi * np.array(v)\n",
        "\n",
        "  return theta, phi\n",
        "    \n",
        "hemisphere = True\n",
        "def get_all_camera_rays(N_cameras, camera_dist, H, W, focal):\n",
        "  theta, phi = nice_uniform_spherical(N_cameras, hemisphere)\n",
        "\n",
        "  camera_x_vec = np.sin(theta) * np.cos(phi)\n",
        "  camera_y_vec = np.sin(theta) * np.sin(phi)\n",
        "  camera_z_vec = np.cos(theta)\n",
        "\n",
        "  rays_o_vec = []\n",
        "  rays_d_vec = []\n",
        "  cameras = []\n",
        "  for i in range(N_cameras):\n",
        "    camera = np.eye(4)\n",
        "    camera[0, 3] = camera_x_vec[i] * camera_dist\n",
        "    camera[1, 3] = camera_y_vec[i] * camera_dist\n",
        "    camera[2, 3] = camera_z_vec[i] * camera_dist\n",
        "\n",
        "    zdir = np.array([camera_x_vec[i], camera_y_vec[i], camera_z_vec[i]])\n",
        "    zdir /= np.linalg.norm(zdir)\n",
        "\n",
        "    ydir = np.array([0.0, 0.0, 1.0])\n",
        "    ydir -= zdir * zdir.dot(ydir)\n",
        "    ydir[0] += 1e-10  # make sure that cameras pointing straight down/up have a defined ydir\n",
        "    ydir /= np.linalg.norm(ydir)\n",
        "\n",
        "    xdir = np.cross(ydir, zdir)\n",
        "\n",
        "\n",
        "    camera[:3, 0] = xdir\n",
        "    camera[:3, 1] = ydir\n",
        "    camera[:3, 2] = zdir\n",
        "\n",
        "    cameras.append(camera)\n",
        "\n",
        "    rays_o, rays_d = get_rays(H, W, focal, camera)\n",
        "\n",
        "    rays_o_vec.append(rays_o)\n",
        "    rays_d_vec.append(rays_d)\n",
        "\n",
        "  rays_o_vec = jnp.stack(rays_o_vec, 0)\n",
        "  rays_d_vec = jnp.stack(rays_d_vec, 0)\n",
        "\n",
        "  return rays_o_vec, rays_d_vec\n",
        "    \n",
        "\n",
        "def plot_cameras(rays_o_vec, rays_d_vec):\n",
        "  fig = go.Figure(data=[go.Scatter3d(\n",
        "    x=rays_o_vec[:, 0, 0, 0].flatten(),\n",
        "    y=rays_o_vec[:, 0, 0, 1].flatten(),\n",
        "    z=rays_o_vec[:, 0, 0, 2].flatten(),\n",
        "    mode='markers',\n",
        "    text=[f'Camera {str(i)}' for i in range(rays_o_vec.shape[0])],\n",
        "    marker=dict(\n",
        "      size=6,\n",
        "      color=[float(i)/rays_o_vec.shape[0] for i in range(rays_o_vec.shape[0])], # set color to an array/list of desired values\n",
        "      colorscale='Viridis',   # choose a colorscale\n",
        "      opacity=0.8\n",
        "    )\n",
        "  )])\n",
        "\n",
        "  ll = 0.1 * np.linalg.norm(rays_o_vec[0, 0, 0, :])  # line length\n",
        "\n",
        "  inds = [(0, 0), (H-1, 0), (H-1, W-1), (0, W-1)]\n",
        "  x = []\n",
        "  y = []\n",
        "  z = []\n",
        "  for ind in inds:\n",
        "    i, j = ind\n",
        "    xyz0 = rays_o_vec[:, i, j, :]\n",
        "    xyz1 = rays_o_vec[:, i, j, :] + ll * rays_d_vec[:, i, j, :]\n",
        "\n",
        "    #x = [xyz0[i, 0], xyz1[i, 0], None for i in range(100)]\n",
        "    x += [item for triplet in zip(xyz0[:, 0], xyz1[:, 0], [None]*N_cameras) for item in triplet]\n",
        "    y += [item for triplet in zip(xyz0[:, 1], xyz1[:, 1], [None]*N_cameras) for item in triplet]\n",
        "    z += [item for triplet in zip(xyz0[:, 2], xyz1[:, 2], [None]*N_cameras) for item in triplet]\n",
        "\n",
        "  for ind, indnext in zip(inds, inds[1:] + inds[:1]):\n",
        "    i, j = ind\n",
        "    inext, jnext = indnext\n",
        "    xyz0 = rays_o_vec[:, i, j, :] + ll * rays_d_vec[:, i, j, :]\n",
        "    xyz1 = rays_o_vec[:, i, j, :] + ll * rays_d_vec[:, inext, jnext, :]\n",
        "\n",
        "    #x = [xyz0[i, 0], xyz1[i, 0], None for i in range(100)]\n",
        "    x += [item for triplet in zip(xyz0[:, 0], xyz1[:, 0], [None]*N_cameras) for item in triplet]\n",
        "    y += [item for triplet in zip(xyz0[:, 1], xyz1[:, 1], [None]*N_cameras) for item in triplet]\n",
        "    z += [item for triplet in zip(xyz0[:, 2], xyz1[:, 2], [None]*N_cameras) for item in triplet]\n",
        "  \n",
        "  fig.add_trace(go.Scatter3d(\n",
        "    x=x, y=y, z=z,\n",
        "    mode='lines',\n",
        "    line=dict(\n",
        "      color='darkblue',\n",
        "      width=2\n",
        "    )\n",
        "  ))\n",
        "\n",
        "  # tight layout\n",
        "  fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))\n",
        "  #fig.show()\n",
        "  \n",
        "  return fig\n",
        "    \n",
        "    \n",
        "def plot_object(texture, shape, fig=None):\n",
        "    N_theta = 100\n",
        "    N_phi = 100\n",
        "\n",
        "    if shape == 'sphere':\n",
        "        theta, phi = np.meshgrid(np.linspace(0.0, np.pi, N_theta), np.linspace(0.0, 2*np.pi, N_phi))\n",
        "        r = sphere_rad\n",
        "    else:\n",
        "        raise ValueError(f'Shape {shape} not implemented.')\n",
        "\n",
        "    x = r * np.sin(theta) * np.cos(phi)\n",
        "    y = r * np.sin(theta) * np.sin(phi)\n",
        "    z = r * np.cos(theta)\n",
        "    xyz = np.stack([x, y, z], -1).reshape(-1, 3)\n",
        "    clr, _, _ = get_albedo_alpha_normals(xyz, sphere_xyz, rad=sphere_rad, texture=texture, shape=shape)\n",
        "\n",
        "    #facecolor=clr.reshape(x.shape[0], x.shape[1], 3)\n",
        "\n",
        "    #surfacecolor = [f'rgb({r}, {g}, {b})' for r, g, b in clr[0, :, :]]\n",
        "    #surfacecolor = ['rgb(0, 0, 0)' for i in range(clr.shape[1])]\n",
        "    colorscale = [[ind, f'rgb({r}, {g}, {b})'] for ind, (r, g, b) in enumerate(set([tuple(c) for c in clr[0, :, :]]))]\n",
        "    surfacecolor = [[t[1] for t in colorscale].index(f'rgb({r}, {g}, {b})') for r, g, b in clr[0, :, :]]\n",
        "\n",
        "    surface = go.Surface(x=x, y=y, z=z,\n",
        "                         colorscale=colorscale,\n",
        "                         surfacecolor=np.array(surfacecolor).reshape(N_theta, N_phi),\n",
        "                         showscale=False,\n",
        "                         cmin=0.0, cmax=1.0)\n",
        "    \n",
        "    # If fig is not specified, create new figure. Else, append.\n",
        "    if fig is None:\n",
        "      fig = go.Figure(data=[surface])\n",
        "      fig.show()\n",
        "    else:\n",
        "      fig.add_trace(surface)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bgjzJXBc-p65"
      },
      "outputs": [],
      "source": [
        "rays_o_vec, rays_d_vec = get_all_camera_rays(200, 1.0, 128, 128, 120.0)\n",
        "plot_cameras(rays_o_vec, rays_d_vec)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lsBMjzCg65dk"
      },
      "outputs": [],
      "source": [
        "def get_gt_materials_and_normals(xyz, oxyz, rad, texture='white', shape='sphere', shading='lambertian'):\n",
        "  \"\"\"\n",
        "  xyz:  shape [N, 3]\n",
        "  oxyz: shape [1, 3]\n",
        "  rad:  float\n",
        "  \n",
        "  returns: materials: dictionary with entries of shape [N, 3]\n",
        "           normals: shape [N, 3]\n",
        "  \"\"\"\n",
        "  assert shape in ['sphere']\n",
        "  x = xyz[:, 0] - oxyz[:, 0]\n",
        "  y = xyz[:, 1] - oxyz[:, 1]\n",
        "  z = xyz[:, 2] - oxyz[:, 2]\n",
        "  \n",
        "  x = x[:, None]\n",
        "  y = y[:, None]\n",
        "  z = z[:, None]\n",
        "  \n",
        "  if shape == 'sphere':\n",
        "    normals = jnp.concatenate([x, y, z], axis=-1)\n",
        "    normals /= jnp.linalg.norm(normals, axis=-1, keepdims=True)\n",
        "\n",
        "  if texture == 'white':\n",
        "    albedo = jnp.ones((x.shape[0], 3), dtype=jnp.float32)\n",
        "\n",
        "  elif texture == 'checkerboard':\n",
        "    theta = jnp.arctan2(jnp.sqrt(x ** 2 + y ** 2), z)\n",
        "    phi   = jnp.arctan2(y, x)\n",
        "\n",
        "    N_phi_stripes = 8\n",
        "    N_theta_stripes = 10\n",
        "    phi_mult   = jnp.where(jnp.remainder(N_phi_stripes   * (phi   + jnp.pi) / (2 * jnp.pi), 2) \u003c 1, 1.0, -1.0)\n",
        "    theta_mult = jnp.where(jnp.remainder(N_theta_stripes * theta / jnp.pi, 2) \u003c 1, 1.0, -1.0)\n",
        "    \n",
        "    clr1 = jnp.array([1.0, 1.0, 1.0])[None, :]\n",
        "    clr2 = jnp.array([1.0, 0.3, 0.3])[None, :]\n",
        "    albedo = jnp.where(phi_mult * theta_mult \u003e 0.0, clr1, clr2)\n",
        "      \n",
        "  elif texture == 'gradient_checkerboard':\n",
        "    theta = jnp.arctan2(jnp.sqrt(x ** 2 + y ** 2), z)\n",
        "    phi   = jnp.arctan2(y, x)\n",
        "\n",
        "    N_phi_stripes = 8\n",
        "    N_theta_stripes = 10\n",
        "    # triangle wave\n",
        "    triangle = lambda x, num: 2.0*num*jnp.abs(jnp.remainder(x - num/8.0, 2.0/num) - 1.0/num) - 1.0\n",
        "    \n",
        "    # Triangle wave in [-1, 1] with period 2/number_of_stripes\n",
        "    phi_mult   = triangle((phi + jnp.pi)/(2.0*jnp.pi), N_phi_stripes)\n",
        "    theta_mult = triangle(theta/jnp.pi, N_theta_stripes)\n",
        "    \n",
        "    mix = phi_mult * theta_mult * 0.5 + 0.5 # map to [0, 1]\n",
        "    \n",
        "    clr1 = jnp.array([1.0, 1.0, 1.0])[None, :]\n",
        "    clr2 = jnp.array([1.0, 0.3, 0.3])[None, :]\n",
        "    albedo = mix * clr1 + (1.0 - mix) * clr2\n",
        "  elif texture == 'noisy_checkerboard':\n",
        "    theta = jnp.arctan2(jnp.sqrt(x ** 2 + y ** 2), z)\n",
        "    phi   = jnp.arctan2(y, x)\n",
        "\n",
        "    N_phi_stripes = 30\n",
        "    N_theta_stripes = 25\n",
        "    phi_ind   = jnp.int32(jnp.floor(N_phi_stripes   * (phi   + jnp.pi) / (2 * jnp.pi)))  # [N, 1]\n",
        "    theta_ind = jnp.int32(jnp.floor(N_theta_stripes * theta / jnp.pi))                   # [N, 1]\n",
        "    \n",
        "    ind = phi_ind * N_theta_stripes + theta_ind\n",
        "    \n",
        "    rng = jax.random.PRNGKey(0)\n",
        "    coeffs = jax.random.uniform(rng, (N_theta_stripes * (N_phi_stripes + 1) + 1,), dtype=jnp.float32)\n",
        "    \n",
        "    mix = coeffs[ind]\n",
        "    \n",
        "    clr1 = jnp.array([1.0, 1.0, 1.0])[None, :]\n",
        "    #clr2 = jnp.array([1.0, 0.3, 0.3])[None, :]\n",
        "    clr2 = jnp.array([0.3, 1.0, 0.3])[None, :]\n",
        "    \n",
        "    albedo = mix * clr1 + (1.0 - mix) * clr2\n",
        "      \n",
        "  else:\n",
        "    raise ValueError(f'Texture {texture} not implemented.')\n",
        "      \n",
        "  materials = {}\n",
        "  materials['albedo'] = albedo\n",
        "  if shading == 'phong':\n",
        "    materials['specular_albedo'] = jnp.ones_like(albedo) * 3.0 #10.0\n",
        "    #materials['specular_exponent'] = jnp.ones_like(albedo) * 100.0\n",
        "    materials['specular_exponent'] = jnp.where((jnp.abs(x) \u003c 0.15 * rad) | (jnp.abs(y) \u003c 0.15 * rad), 20.0, 100.0)\n",
        "  elif shading == 'blinnphong':\n",
        "    materials['specular_albedo'] = jnp.ones_like(albedo) * 5.0 #10.0\n",
        "    #materials['specular_exponent'] = jnp.ones_like(albedo) * 300.0\n",
        "    materials['specular_exponent'] = jnp.where((jnp.abs(x) \u003c 0.15 * rad) | (jnp.abs(y) \u003c 0.15 * rad), 50.0, 300.0)\n",
        "\n",
        "  return materials, normals\n",
        "\n",
        "\n",
        "def get_geometry_and_materials(rays_o, rays_d, oxyz, rad, texture, shape, shading):\n",
        "  # Find point on surface\n",
        "  d_norm_sq = (rays_d ** 2).sum(-1, keepdims=True)\n",
        "  o_norm_sq = (rays_o ** 2).sum(-1, keepdims=True)\n",
        "  d_dot_o = (rays_o * rays_d).sum(-1, keepdims=True)\n",
        "  disc = d_norm_sq * (rad ** 2  - o_norm_sq) + d_dot_o ** 2\n",
        "  alpha = jnp.float32(disc \u003e 0)\n",
        "  t_surface = jnp.where(disc \u003e 0, - jnp.sqrt(disc) - d_dot_o, jnp.inf)  # [H, W, 1]\n",
        "\n",
        "  pts = rays_o + rays_d * t_surface\n",
        "\n",
        "  materials, normals = get_gt_materials_and_normals(pts, oxyz, rad, texture, shape, shading)\n",
        "  return materials, normals, t_surface, alpha\n",
        "\n",
        "\n",
        "def render_pixel(normal, lobe, envmap, mask):\n",
        "  masked_envmap = envmap * mask[:, :, None]\n",
        "  return (masked_envmap * lobe * jnp.sin(omega_theta).reshape(envmap_H, envmap_W, 1)).sum(0).sum(0) * dtheta_dphi / jnp.pi\n",
        "\n",
        "\n",
        "def render_gt(envmap, mask, materials, normals, rays_d, alpha, oxyz, rad, texture='white', shape='sphere', shading='lambertian'):\n",
        "  \"\"\"\n",
        "  envmap:     shape [h, w, 3]\n",
        "  mask:       shape [h, w]\n",
        "  materials:  dictionary with entries of shape [N, 3]\n",
        "  normals:    shape [N, 3]\n",
        "  rays_d:     shape [N, 3]\n",
        "  alpha:      shape [N, 1]\n",
        "  oxyz:       shape [1, 3]  (shape center)\n",
        "  rad:        float         (shape radius)\n",
        "  \n",
        "  output: rendered colors, shape [N, 3]\n",
        "  \"\"\"\n",
        "  \n",
        "  assert shading in ['lambertian', 'phong', 'blinnphong']\n",
        "\n",
        "  if shading in ['lambertian', 'phong', 'blinnphong']:\n",
        "    # TODO: Feed in only pixels where alpha = 1\n",
        "    lobes = jnp.maximum(0.0, (omega_xyz.reshape(1, envmap_H, envmap_W, 3) * normals[:, None, None, :]).sum(-1, keepdims=True)) * materials['albedo'][:, None, None, :]  # [HW, envmap_H, envmap_W, 3]\n",
        "\n",
        "  if shading == 'blinnphong':\n",
        "    assert 'specular_albedo' in materials.keys()\n",
        "    specular_albedo = materials['specular_albedo'][:, None, None, :]\n",
        "    exponent = materials['specular_exponent'][:, None, None, :]\n",
        "\n",
        "    d_norm_sq = (rays_d ** 2).sum(-1, keepdims=True)\n",
        "    rays_d_norm = -rays_d / jnp.sqrt(d_norm_sq + 1e-10)\n",
        "\n",
        "    halfvectors = omega_xyz.reshape(1, envmap_H, envmap_W, 3) + rays_d_norm[:, None, None, :]\n",
        "    halfvectors /= (jnp.linalg.norm(halfvectors, axis=-1, keepdims=True) + 1e-10)  # [N, envmap_H, envmap_W, 3]\n",
        "\n",
        "    lobes += jnp.maximum(0.0, (halfvectors * normals[:, None, None, :]).sum(-1, keepdims=True)) ** exponent * specular_albedo\n",
        "\n",
        "  if shading == 'phong':\n",
        "    assert 'specular_albedo' in materials.keys()\n",
        "    specular_albedo = materials['specular_albedo'][:, None, None, :]\n",
        "    exponent = materials['specular_exponent'][:, None, None, :]\n",
        "\n",
        "    d_norm_sq = (rays_d ** 2).sum(-1, keepdims=True)\n",
        "    rays_d_norm = -rays_d / jnp.sqrt(d_norm_sq + 1e-10)  # [N, 3]\n",
        "\n",
        "    refdirs = 2.0 * (normals * rays_d_norm).sum(-1, keepdims=True) * normals - rays_d_norm   # [N, 3]\n",
        "    refdirs = refdirs[:, None, None, :]\n",
        "\n",
        "    # No need to normalize because ||n|| = 1 and ||d|| = 1, so ||2(n.d)n - d|| = 1.\n",
        "    print(\"Not normalizing here (because unnecessary, at least theoretically).\")\n",
        "    #refdirs /= (jnp.linalg.norm(refdirs, axis=-1, keepdims=True) + 1e-10)  # [N, HW, envmap_H, envmap_W, 3]\n",
        "\n",
        "    lobes += jnp.maximum(0.0, (refdirs * omega_xyz.reshape(1, envmap_H, envmap_W, 3)).sum(-1, keepdims=True)) ** exponent * specular_albedo\n",
        "     \n",
        "  colors = jax.vmap(render_pixel, in_axes=(0, 0, None, None))(normals, lobes, envmap, mask)     \n",
        "  \n",
        "  return colors * alpha"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JRkrgc1TjGyz"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "N_cameras = 2 ** 10\n",
        "focal = 200.0\n",
        "sphere_rad = 0.45*2\n",
        "sphere_xyz = jnp.array([0.0, 0.0, 0.0])[None, :]\n",
        "camera_dist = 3.0\n",
        "near = 2.0\n",
        "far  = 4.0\n",
        "#texture = 'white'\n",
        "texture = 'white'  # 'white', 'checkerboard', 'gradient_checkerboard', 'noisy_checkerboard'\n",
        "#texture = 'gradient_checkerboard'\n",
        "shape = 'sphere'  # 'sphere'\n",
        "shading = 'lambertian'  # 'lambertian', 'phong', 'blinnphong'\n",
        "\n",
        "if shading == 'lambertian':\n",
        "  exposure = 1.0 / 2.0596845\n",
        "elif shading == 'blinnphong':\n",
        "  exposure = 1.0 / 3.865572055570882 #/ 2.7384105199836 #6.199116\n",
        "elif shading == 'phong':\n",
        "  exposure = 1.0 / 6.139414\n",
        "\n",
        "H = W = 128\n",
        "\n",
        "rays_o_vec, rays_d_vec = get_all_camera_rays(N_cameras, camera_dist, H, W, focal)\n",
        "\n",
        "# Make masks\n",
        "mask_shape = 'three_circles'\n",
        "if mask_shape == 'circle' or mask_shape == 'two_circles' or mask_shape == 'three_circles':\n",
        "  masks_gt = jnp.sum(-omega_xyz[None, :, :] * rays_d_vec[:, H//2, W//2, :][:, None, :], axis=-1) \u003c 0.8  # Occluder is aligned with the camera\n",
        "if mask_shape == 'two_circles' or mask_shape == 'three_circles':\n",
        "  camera_dirs = rays_d_vec[:, H//2, W//2, :]\n",
        "  up_dirs = rays_d_vec[:, H//2+1, W//2, :] - camera_dirs\n",
        "  up_dirs = up_dirs / jnp.linalg.norm(up_dirs, axis=-1, keepdims=True)\n",
        "  rot_dirs = jnp.cross(up_dirs, camera_dirs, axis=-1) # direction of rotation is camera direction cross up direction\n",
        "\n",
        "  rotation_angle = jnp.pi / 7\n",
        "  dir2 = camera_dirs * jnp.cos(rotation_angle) + jnp.cross(rot_dirs, camera_dirs, axis=-1) * jnp.sin(rotation_angle) + rot_dirs * (rot_dirs * camera_dirs).sum(-1, keepdims=True) * (1 - jnp.cos(rotation_angle))\n",
        "  masks_gt = jnp.minimum(masks_gt, jnp.sum(-omega_xyz[None, :, :] * dir2[:, None, :], axis=-1) \u003c 0.8)  # Occluder is aligned with the camera\n",
        "if mask_shape == 'three_circles':\n",
        "  camera_dirs = rays_d_vec[:, H//2, W//2, :]\n",
        "  up_dirs = rays_d_vec[:, H//2+1, W//2, :] - camera_dirs\n",
        "  up_dirs = up_dirs / jnp.linalg.norm(up_dirs, axis=-1, keepdims=True)\n",
        "  rot_dirs = up_dirs # direction of rotation is up direction\n",
        "\n",
        "  rotation_angle = jnp.pi / 6\n",
        "  dir2 = camera_dirs * jnp.cos(rotation_angle) + jnp.cross(rot_dirs, camera_dirs, axis=-1) * jnp.sin(rotation_angle) + rot_dirs * (rot_dirs * camera_dirs).sum(-1, keepdims=True) * (1 - jnp.cos(rotation_angle))\n",
        "  masks_gt = jnp.minimum(masks_gt, jnp.sum(-omega_xyz[None, :, :] * dir2[:, None, :], axis=-1) \u003c 0.8)  # Occluder is aligned with the camera\n",
        "\n",
        "masks_gt_old = jnp.float32(masks_gt).reshape(-1, envmap_H, envmap_W)\n",
        "\n",
        "plt.imshow(linear_to_srgb(exposure * envmap_gt))\n",
        "\"\"\";"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3586t4M1EONj"
      },
      "outputs": [],
      "source": [
        "N_cameras = 2 ** 9\n",
        "focal = 200.0\n",
        "sphere_rad = 0.45*2\n",
        "sphere_xyz = jnp.array([0.0, 0.0, 0.0])[None, :]\n",
        "camera_dist = 3.0\n",
        "near = 2.0\n",
        "far  = 4.0\n",
        "#texture = 'white'\n",
        "texture = 'checkerboard'  # 'white', 'checkerboard', 'gradient_checkerboard', 'noisy_checkerboard'\n",
        "#texture = 'gradient_checkerboard'\n",
        "shape = 'sphere'  # 'sphere'\n",
        "shading = 'lambertian'  # 'lambertian', 'phong', 'blinnphong'\n",
        "\n",
        "if shading == 'lambertian':\n",
        "  exposure = 1.0 / 2.0596845 / 1.0031725 / 1.0011\n",
        "elif shading == 'blinnphong':\n",
        "  #exposure = 1.0 / 3.8919785365924775 #/ 2.7384105199836 #6.199116\n",
        "  exposure = 1.0 / 7.090526 / 3.8919785365924775 / 0.20300668\n",
        "elif shading == 'phong':\n",
        "  exposure = 1.0 / 6.139414\n",
        "\n",
        "H = W = 128\n",
        "\n",
        "rays_o_vec, rays_d_vec = get_all_camera_rays(N_cameras, camera_dist, H, W, focal)\n",
        "\n",
        "# Make masks\n",
        "mask_shape = 'three_circles'\n",
        "if mask_shape == 'circle' or mask_shape == 'two_circles' or mask_shape == 'three_circles':\n",
        "  sdfs_gt = 0.8 - jnp.sum(-omega_xyz[None, :, :] * rays_d_vec[:, H//2, W//2, :][:, None, :], axis=-1)\n",
        "if mask_shape == 'two_circles' or mask_shape == 'three_circles':\n",
        "  camera_dirs = rays_d_vec[:, H//2, W//2, :]\n",
        "  up_dirs = rays_d_vec[:, H//2+1, W//2, :] - camera_dirs\n",
        "  up_dirs = up_dirs / jnp.linalg.norm(up_dirs, axis=-1, keepdims=True)\n",
        "  rot_dirs = jnp.cross(up_dirs, camera_dirs, axis=-1) # direction of rotation is camera direction cross up direction\n",
        "\n",
        "  rotation_angle = jnp.pi / 7\n",
        "  dir2 = camera_dirs * jnp.cos(rotation_angle) + jnp.cross(rot_dirs, camera_dirs, axis=-1) * jnp.sin(rotation_angle) + rot_dirs * (rot_dirs * camera_dirs).sum(-1, keepdims=True) * (1 - jnp.cos(rotation_angle))\n",
        "  \n",
        "  sdfs_gt = jnp.minimum(sdfs_gt, 0.8 - jnp.sum(-omega_xyz[None, :, :] * dir2[:, None, :], axis=-1))  # Occluder is aligned with the camera\n",
        "if mask_shape == 'three_circles':\n",
        "  camera_dirs = rays_d_vec[:, H//2, W//2, :]\n",
        "  up_dirs = rays_d_vec[:, H//2+1, W//2, :] - camera_dirs\n",
        "  up_dirs = up_dirs / jnp.linalg.norm(up_dirs, axis=-1, keepdims=True)\n",
        "  rot_dirs = up_dirs # direction of rotation is up direction\n",
        "\n",
        "  rotation_angle = jnp.pi / 6\n",
        "  dir2 = camera_dirs * jnp.cos(rotation_angle) + jnp.cross(rot_dirs, camera_dirs, axis=-1) * jnp.sin(rotation_angle) + rot_dirs * (rot_dirs * camera_dirs).sum(-1, keepdims=True) * (1 - jnp.cos(rotation_angle))\n",
        "  sdfs_gt = jnp.minimum(sdfs_gt, 0.8 - jnp.sum(-omega_xyz[None, :, :] * dir2[:, None, :], axis=-1))  # Occluder is aligned with the camera\n",
        "\n",
        "sdfs_gt = sdfs_gt.reshape(-1, envmap_H, envmap_W)\n",
        "masks_gt = jnp.float32(sdfs_gt \u003e 0.0)\n",
        "\n",
        "plt.imshow(linear_to_srgb(exposure * envmap_gt))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4edfnDErfTpr"
      },
      "outputs": [],
      "source": [
        "# Render dataset\n",
        "num_devices = jax.local_device_count()\n",
        "imgs_gt = []\n",
        "\n",
        "get_stuff = lambda o, d: get_geometry_and_materials(o, d, sphere_xyz, sphere_rad, texture, shape, shading)\n",
        "materials_gt, normals_gt, t_surface_gt, alpha_gt = jax.vmap(get_stuff)(\n",
        "    rays_o_vec.reshape(-1, H*W, 3),\n",
        "    rays_d_vec.reshape(-1, H*W, 3),\n",
        "    )\n",
        "\n",
        "\n",
        "render_gt_partial = functools.partial(render_gt, texture=texture, shape=shape, shading=shading)\n",
        "for i in range(N_cameras // num_devices):\n",
        "  i0 = i * num_devices\n",
        "  i1 = i0 + num_devices\n",
        "\n",
        "  imgs_gt_ = jax.pmap(render_gt_partial, in_axes=(None, 0, 0, 0, 0, 0, None, None))(  # here pmap is faster\n",
        "      envmap_gt * exposure,\n",
        "      masks_gt[i0:i1],\n",
        "      jax.tree_map(lambda x: x[i0:i1], materials_gt),\n",
        "      normals_gt[i0:i1],\n",
        "      rays_d_vec.reshape(-1, H*W, 3)[i0:i1],\n",
        "      alpha_gt[i0:i1],\n",
        "      sphere_xyz,\n",
        "      sphere_rad,\n",
        "      )\n",
        "  imgs_gt.append(imgs_gt_)\n",
        "\n",
        "imgs_gt = jnp.concatenate(imgs_gt, axis=0)\n",
        "\n",
        "print(jnp.nanmax(imgs_gt))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0TYghQu3HQOz"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=[6, 6])\n",
        "plt.imshow(imgs_gt[100].reshape(H, W, 3))\n",
        "plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "63vtd6PDItfV"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=[6, 6])\n",
        "plt.imshow(imgs_gt[100].reshape(H, W, 3))\n",
        "plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aOuQ2JYm2mNX"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=[6, 6])\n",
        "plt.imshow(imgs_gt[100].reshape(H, W, 3))\n",
        "plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2Tvo2cTdv-UY"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "num_devices = jax.local_device_count()\n",
        "\n",
        "rays_o_r = rays_o_vec.reshape(-1, H*W, 3)\n",
        "rays_d_r = rays_d_vec.reshape(-1, H*W, 3)\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  features: Sequence[int]\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x, y=None):\n",
        "    if y is not None:\n",
        "      x = jnp.concatenate([x, y], axis=-1)\n",
        "    for feat in self.features[:-1]:\n",
        "      x = nn.relu(nn.Dense(feat)(x))\n",
        "    x = nn.Dense(self.features[-1])(x)\n",
        "    return x\n",
        "\n",
        "\n",
        "append_identity = True\n",
        "def posenc(x, L_encoding):\n",
        "  if L_encoding \u003c= 0:\n",
        "    return x\n",
        "  else:\n",
        "    scales = 2**jnp.arange(L_encoding)\n",
        "\n",
        "    shape = x.shape[:-1] + (-1,)\n",
        "    scaled_x = x[..., None, :] * scales[:, None] # [..., L, D]\n",
        "\n",
        "    four_feat = jnp.sin(\n",
        "        jnp.stack([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1)) # [..., L, D, 2]\n",
        "\n",
        "    four_feat = jnp.reshape(four_feat / scales[:, None, None], shape)\n",
        "    print(\"Using Lipschitz posenc\")\n",
        "    if append_identity:\n",
        "      return jnp.concatenate([x] + [four_feat], axis=-1)\n",
        "    else:\n",
        "      return four_feat\n",
        "\n",
        "\n",
        "\n",
        "def params_to_materials(params_materials, pts):\n",
        "  mlp_res = material_mlp.apply(params_materials, posenc(pts, L_encoding_materials))\n",
        "  materials = {}\n",
        "  materials['albedo'] = jax.nn.sigmoid(mlp_res[..., 0:3])\n",
        "  if shading in ['phong', 'blinnphong']:\n",
        "    materials['specular_albedo'] = 20.0 * jax.nn.sigmoid(mlp_res[..., 3:6])\n",
        "    materials['specular_exponent'] = jax.nn.softplus(mlp_res[..., 6:9])\n",
        "\n",
        "  return materials\n",
        "\n",
        "\n",
        "def get_loss(params_materials, gt, spatial_inds, img_inds, i, rng):\n",
        "  rng, key = jax.random.split(rng)\n",
        "\n",
        "  envmap = envmap_gt\n",
        "\n",
        "  normals = normals_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  t_surface = t_surface_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  alpha = alpha_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  rays_d = rays_d_r[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  rays_o = rays_o_r[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "\n",
        "  #materials = jax.tree_map(lambda x: x[img_inds[:, None], spatial_inds[img_inds], ...], materials_gt)\n",
        "\n",
        "  pts = rays_o + rays_d * t_surface\n",
        "\n",
        "  #materials = params_to_materials(params_materials, jnp.zeros((spatial_batch_size, 3)))\n",
        "  materials = params_to_materials(params_materials, pts)\n",
        "  #materials = {'albedo': materials_gt['albedo'][img_inds[:, None], spatial_inds[img_inds], ...],\n",
        "  #             'specular_albedo': materials['specular_albedo'],\n",
        "  #             'specular_exponent': jnp.zeros_like(rays_d)}\n",
        "  #materials = {'albedo': params_materials[img_inds[:, None], spatial_inds[img_inds], ...],\n",
        "  #             'specular_albedo': materials_gt['specular_albedo'][img_inds[:, None], spatial_inds[img_inds], ...],\n",
        "  #             'specular_exponent': materials_gt['specular_exponent'][img_inds[:, None], spatial_inds[img_inds], ...]}\n",
        "\n",
        "  masks = masks_gt[img_inds]\n",
        "\n",
        "  res = jax.vmap(render_gt_partial, in_axes=(None, 0, 0, 0, 0, 0, None, None))(\n",
        "      envmap,\n",
        "      masks,\n",
        "      materials,\n",
        "      normals,\n",
        "      rays_d,\n",
        "      alpha,\n",
        "      sphere_xyz,\n",
        "      sphere_rad,\n",
        "      )\n",
        "  \n",
        "  diff = gt[img_inds[:, None], spatial_inds[img_inds], :] - res\n",
        "\n",
        "  data_loss = jnp.power((jnp.abs(diff + 1e-10) ** 2).sum(), 1/2.0)\n",
        "  data_loss = data_loss / img_inds.shape[0] / spatial_inds.shape[-1]\n",
        "\n",
        "  loss = data_loss\n",
        "\n",
        "  return loss, (data_loss,)\n",
        "\n",
        "  \n",
        "\n",
        "def params_to_components(params_envmap):\n",
        "  envmap = jax.nn.sigmoid(params_envmap)\n",
        "  return envmap\n",
        "\n",
        "\n",
        "def update_params(i, rng, state_materials, gt, spatial_inds, img_inds):\n",
        "  params_materials = get_params_materials(state_materials)\n",
        "  print(\"Hey\")\n",
        "\n",
        "  (loss, (data_loss,)), g = jax.value_and_grad(get_loss, argnums=(0,), has_aux=True)(params_materials,\n",
        "                                                                                    gt, spatial_inds,\n",
        "                                                                                    img_inds, i, rng)\n",
        "\n",
        "  grad_materials = jax.lax.pmean(g[0], axis_name='batch')\n",
        "  data_loss = jax.lax.pmean(data_loss, axis_name='batch')\n",
        "  loss = jax.lax.pmean(loss, axis_name='batch')\n",
        "\n",
        "  return state_materials, loss, data_loss\n",
        "\n",
        "\n",
        "# Initialize material MLP\n",
        "L_encoding_materials = 0\n",
        "mlp_input_features = 3 + 6 * L_encoding_materials\n",
        "material_mlp = MLP([128]*1 + [3 + 6 * int(shading in ['phong', 'blinnphong'])])\n",
        "\n",
        "params_materials = material_mlp.init(jax.random.PRNGKey(0),\n",
        "                          np.zeros([1, mlp_input_features]))\n",
        "\n",
        "init_lr_materials = 0.0001\n",
        "init_materials, update_materials, get_params_materials = jax.experimental.optimizers.adam(init_lr_materials)\n",
        "state_materials = init_materials(params_materials)\n",
        "\n",
        "\n",
        "np_rng = np.random.default_rng(12345)\n",
        "jax_rng = jax.random.PRNGKey(3948)\n",
        "\n",
        "spatial_batch_size = 1\n",
        "img_batch_size = 8\n",
        "\n",
        "replicated_state_materials = flax.jax_utils.replicate(state_materials)\n",
        "replicated_imgs_gt = flax.jax_utils.replicate(imgs_gt)\n",
        "\n",
        "p_update = jax.pmap(update_params, in_axes=(None, None, 0, 0, 0, 0), axis_name='batch')\n",
        "\n",
        "for iteration in range(1):\n",
        "  # Generate B1 image indices\n",
        "  img_inds = np_rng.choice(imgs_gt.shape[0], size=img_batch_size, replace=False)\n",
        "  # Now generate B2 pixel indices for each image. The total batch size is B1 * B2.\n",
        "\n",
        "  keys = jax.random.split(jax_rng, num=img_batch_size+1)\n",
        "  jax_rng, keys = keys[0], keys[1:]\n",
        "  spatial_inds = jax.vmap(jax.random.choice, in_axes=(0, None, None, None, 0))(keys, H*W, (spatial_batch_size,), False, alpha_gt[img_inds, :, 0])\n",
        "\n",
        "  replicated_state_materials, l1, l2 = p_update(\n",
        "      iteration,\n",
        "      jax_rng,\n",
        "      replicated_state_materials,\n",
        "      replicated_imgs_gt,\n",
        "      spatial_inds.reshape(num_devices, -1, spatial_batch_size),\n",
        "      img_inds.reshape(num_devices, -1)\n",
        "      )\n",
        "  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CLJ7TswqdzUT"
      },
      "outputs": [],
      "source": [
        "import scipy.special\n",
        "N = envmap_H\n",
        "\n",
        "x_vals_full, w_vals_full = scipy.special.roots_legendre(N)\n",
        "x_vals = x_vals_full[N//2:]\n",
        "w_vals = w_vals_full[N//2:]\n",
        "\n",
        "dtype = np.float64\n",
        "\n",
        "ell = jnp.arange(N)[:, None]\n",
        "em  = jnp.arange(N)[None, :]\n",
        "k = em + 1\n",
        "diag_values = jnp.sqrt(jnp.prod(jnp.where(ell \u003e= em + 1, (2.0 * k + 1.0) / 2.0 / k, 1), axis=1) / 4.0 / jnp.pi)\n",
        "a = jnp.where(ell == em, jnp.diag(diag_values), jnp.sqrt((4 * ell ** 2 - 1.0) / (ell ** 2 - em ** 2)))\n",
        "b = -jnp.sqrt((2.0 * ell + 1.0) * ((ell - 1) ** 2 - em ** 2) / (2 * ell - 3) / (ell ** 2 - em ** 2))\n",
        "\n",
        "P = np.zeros((N, N, x_vals.shape[0]), dtype=dtype) + np.nan\n",
        "x = jnp.flip(x_vals)\n",
        "for n in range(N):\n",
        "  for m in range(n, -1, -1): \n",
        "    if m == n:\n",
        "      P[n, m, :] = a[m, m] * (1.0 - x ** 2) ** (abs(m) / 2.0)\n",
        "    elif n == m + 1:\n",
        "      P[n, m, :] = a[n, m] * x * P[m, m, :]\n",
        "    else:\n",
        "      P[n, m, :] = a[n, m] * x * P[n - 1, m, :] + b[n, m] * P[n - 2, m, :]\n",
        "\n",
        "sgn = (-1) ** (jnp.arange(N)[:, None] + jnp.arange(N)[None, :])\n",
        "P_stacked = jnp.concatenate([P, sgn[:, :, None] * jnp.flip(P, axis=2)], axis=-1) # this is P for all theta values from 0 to pi (not only 0 to pi/2)\n",
        "\n",
        "def sht(envmap):\n",
        "  \"\"\"\n",
        "  Equations 8-11 of:\n",
        "  https://arxiv.org/pdf/1202.6522.pdf\n",
        "\n",
        "  Also using the \"mirror symmetry\" property and the fact that since the signal\n",
        "  is real, we can reuse the same polynomial but flipping the negative FFT.\n",
        "  \"\"\"\n",
        "\n",
        "  fm = jnp.fft.fft(envmap, axis=1)[:, :N].T  # FFT w.r.t. azimuth: this is a function of m (row index) and theta (column index)\n",
        "                                             # Only need the first N coefficients (I think?). So this can be made more efficient:\n",
        "                                             # TODO: replace jnp.fft.fft with jnp.fft.rfft\n",
        "\n",
        "  w = N // 2\n",
        "\n",
        "  # Flip w values because P values (and everything else involving x_vals and w_vals) is flipped\n",
        "  summed_fft = fm[None, :, :w] + sgn[:, :, None] * jnp.flip(fm[None, :, w:], axis=2)\n",
        "  folded = summed_fft * jnp.flip(w_vals)[None, None, :] * P  # TODO: Reindex P as [n, x_ind, m] so that we don't have to transpose fm\n",
        "  coeffs = folded.sum(-1)\n",
        "  return coeffs / 8.0 / jnp.pi\n",
        "\n",
        "\n",
        "def isht(coeffs):\n",
        "  \"\"\"\n",
        "  Equation 6 of:\n",
        "  https://arxiv.org/pdf/1202.6522.pdf\n",
        "\n",
        "  Also using:\n",
        "  1. For theta \u003e pi/2, P_n^m(cos(theta)) = (-1)^{n+m} * P_n^m(cos(pi-theta)).\n",
        "  2. For m \u003c 0 we have (from Equation 9) that the FFT is simply reflected (about m) and conjugated\n",
        "  \"\"\"\n",
        "\n",
        "  # Get FFT w.r.t. azimuth by summing across all n \u003e= |m|. This yields the nonnegative m values of the FFT.\n",
        "  positive_fft = jnp.nansum(coeffs[:, :, None] * P_stacked, axis=0) # first dim is m, second dim is cos(theta)\n",
        "  \n",
        "  positive_fft = positive_fft.T  # TODO: Get rid of this by reindexing\n",
        "  # Use symmetry to get the negative m part of the FFT\n",
        "  negative_fft = jnp.conj(jnp.flip(positive_fft[:, 1:], axis=1))\n",
        "\n",
        "  fft = jnp.concatenate([positive_fft, negative_fft], axis=1)  # Uhhh why am I using positive m first in the FFT?\n",
        "\n",
        "  return 16.0 * jnp.pi ** 2 * jnp.real(jnp.fft.ifft(fft, axis=1))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NoLT2X78SDJh"
      },
      "outputs": [],
      "source": [
        "stop running here"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "t-C2MXf7c4X-"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "num_devices = jax.local_device_count()\n",
        "\n",
        "rays_o_r = rays_o_vec.reshape(-1, H*W, 3)\n",
        "rays_d_r = rays_d_vec.reshape(-1, H*W, 3)\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  features: Sequence[int]\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x, y=None):\n",
        "    if y is not None:\n",
        "      x = jnp.concatenate([x, y], axis=-1)\n",
        "    for feat in self.features[:-1]:\n",
        "      x = nn.relu(nn.Dense(feat)(x))\n",
        "    x = nn.Dense(self.features[-1])(x)\n",
        "    return x\n",
        "\n",
        "\n",
        "append_identity = True\n",
        "def posenc(x, L_encoding):\n",
        "  if L_encoding \u003c= 0:\n",
        "    return x\n",
        "  else:\n",
        "    scales = 2**jnp.arange(L_encoding)\n",
        "\n",
        "    shape = x.shape[:-1] + (-1,)\n",
        "    scaled_x = x[..., None, :] * scales[:, None] # [..., L, D]\n",
        "\n",
        "    four_feat = jnp.sin(\n",
        "        jnp.stack([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1)) # [..., L, D, 2]\n",
        "\n",
        "    four_feat = jnp.reshape(four_feat / scales[:, None, None], shape)\n",
        "    print(\"Using Lipschitz posenc\")\n",
        "    if append_identity:\n",
        "      return jnp.concatenate([x] + [four_feat], axis=-1)\n",
        "    else:\n",
        "      return four_feat\n",
        "\n",
        "\n",
        "\n",
        "def params_to_materials(params_materials, pts):\n",
        "  mlp_res = material_mlp.apply(params_materials, posenc(pts, L_encoding_materials))\n",
        "  materials = {}\n",
        "  materials['albedo'] = jax.nn.sigmoid(mlp_res[..., 0:3])\n",
        "  if shading in ['phong', 'blinnphong']:\n",
        "    materials['specular_albedo'] = 20.0 * jax.nn.sigmoid(mlp_res[..., 3:6])\n",
        "    materials['specular_exponent'] = jax.nn.softplus(mlp_res[..., 6:9])\n",
        "\n",
        "  return materials\n",
        "\n",
        "\n",
        "def get_loss(params_materials, gt, spatial_inds, img_inds, i, rng):\n",
        "  rng, key = jax.random.split(rng)\n",
        "\n",
        "  envmap = envmap_gt\n",
        "\n",
        "  normals = normals_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  t_surface = t_surface_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  alpha = alpha_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  rays_d = rays_d_r[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  rays_o = rays_o_r[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "\n",
        "  #materials = jax.tree_map(lambda x: x[img_inds[:, None], spatial_inds[img_inds], ...], materials_gt)\n",
        "\n",
        "  pts = rays_o + rays_d * t_surface\n",
        "\n",
        "  #materials = params_to_materials(params_materials, jnp.zeros((spatial_batch_size, 3)))\n",
        "  materials = params_to_materials(params_materials, pts)\n",
        "  #materials = {'albedo': materials_gt['albedo'][img_inds[:, None], spatial_inds[img_inds], ...],\n",
        "  #             'specular_albedo': materials['specular_albedo'],\n",
        "  #             'specular_exponent': jnp.zeros_like(rays_d)}\n",
        "  #materials = {'albedo': params_materials[img_inds[:, None], spatial_inds[img_inds], ...],\n",
        "  #             'specular_albedo': materials_gt['specular_albedo'][img_inds[:, None], spatial_inds[img_inds], ...],\n",
        "  #             'specular_exponent': materials_gt['specular_exponent'][img_inds[:, None], spatial_inds[img_inds], ...]}\n",
        "\n",
        "  masks = masks_gt[img_inds]\n",
        "\n",
        "  res = jax.vmap(render_gt_partial, in_axes=(None, 0, 0, 0, 0, 0, None, None))(\n",
        "      envmap,\n",
        "      masks,\n",
        "      materials,\n",
        "      normals,\n",
        "      rays_d,\n",
        "      alpha,\n",
        "      sphere_xyz,\n",
        "      sphere_rad,\n",
        "      )\n",
        "  \n",
        "  diff = gt[img_inds[:, None], spatial_inds[img_inds], :] - res\n",
        "\n",
        "  data_loss = jnp.power((jnp.abs(diff + 1e-10) ** 2).sum(), 1/2.0)\n",
        "  data_loss = data_loss / img_inds.shape[0] / spatial_inds.shape[-1]\n",
        "\n",
        "  loss = data_loss\n",
        "\n",
        "  return loss, (data_loss,)\n",
        "\n",
        "  \n",
        "\n",
        "def params_to_components(params_envmap):\n",
        "  envmap = jax.nn.sigmoid(params_envmap)\n",
        "  return envmap\n",
        "\n",
        "\n",
        "def update_params(i, rng, state_materials, gt, spatial_inds, img_inds):\n",
        "  params_materials = get_params_materials(state_materials)\n",
        "  print(\"Hey\")\n",
        "\n",
        "  (loss, (data_loss,)), g = jax.value_and_grad(get_loss, argnums=(0,), has_aux=True)(params_materials,\n",
        "                                                                                    gt, spatial_inds,\n",
        "                                                                                    img_inds, i, rng)\n",
        "\n",
        "  grad_materials = jax.lax.pmean(g[0], axis_name='batch')\n",
        "  data_loss = jax.lax.pmean(data_loss, axis_name='batch')\n",
        "  loss = jax.lax.pmean(loss, axis_name='batch')\n",
        "\n",
        "  return state_materials, loss, data_loss\n",
        "\n",
        "\n",
        "# Initialize material MLP\n",
        "L_encoding_materials = 0\n",
        "mlp_input_features = 3 + 6 * L_encoding_materials\n",
        "material_mlp = MLP([128]*1 + [3 + 6 * int(shading in ['phong', 'blinnphong'])])\n",
        "\n",
        "params_materials = material_mlp.init(jax.random.PRNGKey(0),\n",
        "                          np.zeros([1, mlp_input_features]))\n",
        "\n",
        "init_lr_materials = 0.0001\n",
        "init_materials, update_materials, get_params_materials = jax.experimental.optimizers.adam(init_lr_materials)\n",
        "state_materials = init_materials(params_materials)\n",
        "\n",
        "\n",
        "np_rng = np.random.default_rng(12345)\n",
        "jax_rng = jax.random.PRNGKey(3948)\n",
        "\n",
        "spatial_batch_size = 1\n",
        "img_batch_size = 8\n",
        "\n",
        "replicated_state_materials = flax.jax_utils.replicate(state_materials)\n",
        "replicated_imgs_gt = flax.jax_utils.replicate(imgs_gt)\n",
        "\n",
        "print(\"He\")\n",
        "for iteration in range(1):\n",
        "  print(\"Hyyyy\")\n",
        "  # Generate B1 image indices\n",
        "  img_inds = np_rng.choice(imgs_gt.shape[0], size=img_batch_size, replace=False)\n",
        "  # Now generate B2 pixel indices for each image. The total batch size is B1 * B2.\n",
        "\n",
        "  keys = jax.random.split(jax_rng, num=img_batch_size+1)\n",
        "  jax_rng, keys = keys[0], keys[1:]\n",
        "  spatial_inds = jax.vmap(jax.random.choice, in_axes=(0, None, None, None, 0))(keys, H*W, (spatial_batch_size,), False, alpha_gt[img_inds, :, 0])\n",
        "\n",
        "  replicated_state_materials, l1, l2 = jax.pmap(update_params, in_axes=(None, None, 0, 0, 0, 0), axis_name='batch')(\n",
        "      iteration,\n",
        "      jax_rng,\n",
        "      replicated_state_materials,\n",
        "      replicated_imgs_gt,\n",
        "      spatial_inds.reshape(num_devices, -1, spatial_batch_size),\n",
        "      img_inds.reshape(num_devices, -1)\n",
        "      )\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uKhtJ6iZS4x5"
      },
      "source": [
        "Entropy loss is:\n",
        "$$-m\\log(m)-(1-m)\\log(1-m)$$\n",
        "\n",
        "But because $m = \\sigma(d)$ where $d$ is an SDF, we have that the entropy loss is:\n",
        "$$-\\sigma(d)\\log(\\sigma(d)) - (1-\\sigma(d))\\log(1-\\sigma(d)) \\\\\n",
        "= \\frac{\\log(1+e^{-d})}{1+e^{-d}} + \\frac{e^{-d}[\\log(1+e^{-d}) + d]}{1+e^{-d}} \\\\\n",
        "= \\log(1+e^{-d}) + d\\cdot(1-\\sigma(d)).$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0G89obOsUot_"
      },
      "outputs": [],
      "source": [
        "d = jnp.linspace(-5.0, 5.0, 1000)\n",
        "m = jax.nn.sigmoid(d)\n",
        "\n",
        "#plt.plot(d, -m * jnp.log(m) - (1.0 - m) * jnp.log(1.0 - m))\n",
        "plt.plot(d, jnp.log(1+jnp.exp(-d)) + d * (1 - jax.nn.sigmoid(d)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hpGX-r99WG_P"
      },
      "outputs": [],
      "source": [
        "num_devices = jax.local_device_count()\n",
        "\n",
        "gt_list = ['materials']\n",
        "#gt_list = ['masks', 'materials']\n",
        "#gt_list = ['sdfs', 'materials']\n",
        "#gt_list = ['envmap', 'materials']\n",
        "#gt_list = []\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  features: Sequence[int]\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x, y=None):\n",
        "    if y is not None:\n",
        "      x = jnp.concatenate([x, y], axis=-1)\n",
        "    for feat in self.features[:-1]:\n",
        "      x = nn.relu(nn.Dense(feat)(x))\n",
        "    x = nn.Dense(self.features[-1])(x)\n",
        "    return x\n",
        "\n",
        "def get_pyramid_params(rng, num_pyramids, height, width, pyramid_num_scales, pyramid_resize_scale, global_std, global_bias):\n",
        "  pyramid_params = []\n",
        "  for i in range(pyramid_num_scales):\n",
        "    gsh, gsw = [sz // pyramid_resize_scale ** i for sz in [height, width]]\n",
        "    key, rng = jax.random.split(rng)\n",
        "    features = jax.random.normal(key, (num_pyramids, gsh, gsw)) * global_std + global_bias\n",
        "    pyramid_params.append(features)\n",
        "  return pyramid_params\n",
        "\n",
        "\n",
        "def pyramids_to_imgs(pyramid_params, pyramid_mult, img_inds):\n",
        "  def pyramid_to_img(pyramid_params, pyramid_mult):\n",
        "    acc_val = pyramid_params[-1]\n",
        "    for i, curr_val in enumerate(pyramid_params[-2::-1], start=1):\n",
        "      # upsample\n",
        "      acc_val = jax.image.resize(acc_val * pyramid_mult, shape=curr_val.shape, method='linear')\n",
        "      \n",
        "      # accumulate\n",
        "      acc_val += curr_val\n",
        "    return acc_val\n",
        "\n",
        "  # Select all pyramid parameters at given indices\n",
        "  sub_pyramid_params = jax.tree_map(lambda t: t[img_inds], pyramid_params)\n",
        "  return jax.vmap(pyramid_to_img, in_axes=(0, None))(sub_pyramid_params, pyramid_mult)\n",
        "\n",
        "\n",
        "def grad_norm_spherical(dirs, grad):\n",
        "  \"\"\"\n",
        "  Compute gradient norm restricted to the sphere.\n",
        "  Assume dim 0 is elevation and 1 is azimuth.\n",
        "\n",
        "  dirs: (N, 3) array of directions on the sphere\n",
        "  grad: (N, 3) array of Cartesian gradients of points on dirs\n",
        "  \"\"\"\n",
        "\n",
        "  norm_spherical = grad - (dirs * grad).sum(axis=-1, keepdims=True) * dirs\n",
        "  return jnp.sqrt((norm_spherical ** 2).sum(-1) + 1e-5)\n",
        "\n",
        "\n",
        "def add_fake_quantization_noise(x, num_bits):\n",
        "  m = jnp.nanmin(x)\n",
        "  M = jnp.nanmax(x)\n",
        "  x = (x - m) / (M - m)\n",
        "  #return jnp.float32(jnp.uint8(x * 255.0)) / 255 * (M - m) + m\n",
        "  if num_bits == 8:\n",
        "    return np.float32(np.uint8(x * 255.0)) / 255.0 * (M - m) + m\n",
        "    \n",
        "  d = 0.5 ** num_bits  # This is the distance between quantization bins, and the bins' widths\n",
        "  x = 0.5 * d + (1.0 - d) * x\n",
        "  return (x + (np.random.rand(*x.shape) - 0.5) * d) * (M - m) + m\n",
        "\n",
        "rays_o_r = rays_o_vec.reshape(-1, H*W, 3)\n",
        "rays_d_r = rays_d_vec.reshape(-1, H*W, 3)\n",
        "\n",
        "num_bits = 1024\n",
        "imgs_gt_noisy = add_fake_quantization_noise(imgs_gt, num_bits)\n",
        "for _ in range(10):\n",
        "  print(f\"Using {num_bits} bits\")\n",
        "\n",
        "\n",
        "append_identity = True\n",
        "def posenc(x, L_encoding):\n",
        "  if L_encoding \u003c= 0:\n",
        "    return x\n",
        "  else:\n",
        "    scales = 2**jnp.arange(L_encoding)\n",
        "    #shape = x.shape[:-1] + (-1,)\n",
        "    #scaled_x = jnp.reshape((x[..., None, :] * scales[:, None]), shape)\n",
        "\n",
        "    #four_feat = jnp.sin(\n",
        "    #    jnp.concatenate([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1))\n",
        "    shape = x.shape[:-1] + (-1,)\n",
        "    scaled_x = x[..., None, :] * scales[:, None] # [..., L, D]\n",
        "\n",
        "    four_feat = jnp.sin(\n",
        "        jnp.stack([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1)) # [..., L, D, 2]\n",
        "\n",
        "    four_feat = jnp.reshape(four_feat / scales[:, None, None], shape)\n",
        "    print(\"Using Lipschitz posenc\")\n",
        "    if append_identity:\n",
        "      return jnp.concatenate([x] + [four_feat], axis=-1)\n",
        "    else:\n",
        "      return four_feat\n",
        "\n",
        "\n",
        "def params_to_sdf(params_sdf, img_inds):\n",
        "  if sdf_representation == 'mlp':\n",
        "    sdf, sdf_grad = jax.vmap(jax.value_and_grad(lambda x, y: sdf_mlp.apply(params_sdf, x, y)[0]))(\n",
        "        posenc(omega_xyz[None, :, :], L_encoding_sdf).repeat(img_inds.shape[0], 0).reshape(-1, mlp_input_features),\n",
        "        rays_o_vec[img_inds, 0, 0, :][:, None, :].repeat(envmap_H * envmap_W, 1).reshape(-1, 3)\n",
        "        )\n",
        "  elif sdf_representation == 'grid':\n",
        "    sdf = params_sdf[img_inds[:, None], ...]\n",
        "  else:\n",
        "    sdf = pyramids_to_imgs(params_sdf, pyramid_mult, img_inds)\n",
        "  return sdf\n",
        "\n",
        "\n",
        "def sdf_to_mask(x, width, curve='sigmoid'):\n",
        "  if curve == 'sigmoid':\n",
        "    return jax.nn.sigmoid(x * width)\n",
        "  elif curve == 'laplace_cdf':\n",
        "    return 0.5 + 0.5 * jnp.sign(x) * (1.0 - jnp.exp(-jnp.abs(x) * width))\n",
        "  else:\n",
        "    raise NotImplementedError('Only sigmoid and laplace_cdf for now.')\n",
        "  \n",
        "def params_to_materials(params_materials, pts):\n",
        "  mlp_res = material_mlp.apply(params_materials, posenc(pts, L_encoding_materials))\n",
        "  materials = {}\n",
        "  materials['albedo'] = jax.nn.sigmoid(mlp_res[..., 0:3])\n",
        "  if shading in ['phong', 'blinnphong']:\n",
        "    #materials['specular_albedo'] = 20.0 * jax.nn.sigmoid(mlp_res[..., 3:6])\n",
        "    materials['specular_albedo'] = jax.nn.softplus(mlp_res[..., 3:6] - 3.0)\n",
        "    materials['specular_exponent'] = jax.nn.softplus(mlp_res[..., 6:9])\n",
        "\n",
        "  return materials\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def get_loss(params_envmap, params_sdf, params_materials, gt, spatial_inds, img_inds, i, rng):\n",
        "  rng, key = jax.random.split(rng)\n",
        "\n",
        "  if 'envmap' in gt_list:\n",
        "    envmap = envmap_gt * exposure\n",
        "  else:\n",
        "    envmap = params_to_envmap(params_envmap)\n",
        "\n",
        "  normals = normals_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  t_surface = t_surface_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  alpha = alpha_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  rays_d = rays_d_r[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  rays_o = rays_o_r[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "\n",
        "  if 'materials' in gt_list:\n",
        "    materials = jax.tree_map(lambda x: x[img_inds[:, None], spatial_inds[img_inds], ...], materials_gt)\n",
        "  else:\n",
        "    pts = rays_o + rays_d * t_surface\n",
        "    materials = params_to_materials(params_materials, pts)\n",
        "\n",
        "  sdf = params_to_sdf(params_sdf, img_inds)\n",
        "  sdf = sdf.reshape(img_inds.shape[0], envmap_H * envmap_W)\n",
        "  #mask_width = 200.0 * (i / num_iters) + 10.0\n",
        "  #masks = sdf_to_mask(sdf.reshape(-1, envmap_H, envmap_W), mask_width, 'sigmoid')\n",
        "\n",
        "  sdf_curve = 'sigmoid'\n",
        "\n",
        "  if 'masks' in gt_list or 'sdfs' in gt_list:\n",
        "    if 'masks' in gt_list:\n",
        "      masks = masks_gt[img_inds]\n",
        "    else:\n",
        "      masks = sdf_to_mask(sdfs_gt[img_inds], mask_width, sdf_curve)\n",
        "    eikonal_loss = 0.0\n",
        "    length_loss = 0.0\n",
        "    mask_area_loss = 0.0\n",
        "  else:\n",
        "    #mask_width = 10.0 #200.0 * (i / num_iters) + 10.0\n",
        "    #mask_width = 20.0 * (i / num_iters) + 10.0\n",
        "    #mask_width = 0.1\n",
        "    #mask_width = 10.0 #150.0 * (i / num_iters) + 6.0\n",
        "    #mask_width = 0.1 * jnp.exp(jnp.log(100) * i / num_iters)\n",
        "    if straight_through_mode == 'hard':\n",
        "      masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "      masks = (masks + jax.lax.stop_gradient(jnp.float32(sdf \u003e 0.0) - masks))\n",
        "    elif straight_through_mode == 'soft':\n",
        "      soft_masks = sdf_to_mask(sdf, 0.1, sdf_curve)\n",
        "      hard_masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "      # Define masks with value of `hard_masks` but gradients of `soft_masks`\n",
        "      masks = (soft_masks + jax.lax.stop_gradient(hard_masks - soft_masks))\n",
        "    elif straight_through_mode == 'none':\n",
        "      masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "\n",
        "\n",
        "\n",
        "    if sdf_representation == 'mlp':\n",
        "      # TODO: Only compute grad w.r.t. x, not w.r.t. other posenc components\n",
        "      sdf_grad = sdf_grad[..., :3].reshape(img_inds.shape[0], envmap_H * envmap_W, 3)\n",
        "      sdf_grad_norm = grad_norm_spherical(omega_xyz, sdf_grad.reshape(img_inds.shape[0], envmap_H * envmap_W, 3))\n",
        "      eikonal_loss = (jnp.sin(omega_theta) * (sdf_grad_norm - 1) ** 2).sum() * dtheta_dphi\n",
        "    else:\n",
        "      eikonal_loss = 0.0\n",
        "\n",
        "    # Compute entropy assuming mask = sigmoid(sdf)\n",
        "    entropy_loss = jax.nn.softplus(-sdf) + sdf * (1.0 - jax.nn.sigmoid(sdf))\n",
        "    mask_area_loss = 1.0 - masks  # Try to make the occluder as small as possible\n",
        "    #delta_sdf = jax.vmap(jax.grad(lambda x: sdf_to_mask(x, mask_width, sdf_curve)))(sdf.flatten()).reshape(img_inds.shape[0], envmap_H * envmap_W)\n",
        "\n",
        "    #length_loss = (jnp.sin(omega_theta) * sdf_grad_norm * delta_sdf).sum() * dtheta_dphi\n",
        "\n",
        "  res = jax.vmap(render_gt_partial, in_axes=(None, 0, 0, 0, 0, 0, None, None))(\n",
        "      envmap,\n",
        "      masks.reshape(img_inds.shape[0], envmap_H, envmap_W),\n",
        "      materials,\n",
        "      normals,\n",
        "      rays_d,\n",
        "      alpha,\n",
        "      sphere_xyz,\n",
        "      sphere_rad,\n",
        "      )\n",
        "\n",
        "  #diff = gt[img_inds[:, None], spatial_inds[img_inds], :] - res\n",
        "  diff = linear_to_srgb(gt[img_inds[:, None], spatial_inds[img_inds], :]) - linear_to_srgb(res)\n",
        "  #loss_per_element = (diff ** 2).sum(-1)\n",
        "  #loss_per_element = jnp.abs(diff).sum(-1)\n",
        "  if False:\n",
        "    p = 2 - 1.5 * i / num_iters\n",
        "    data_loss = jnp.power((jnp.abs(diff + 1e-10) ** p).sum(), 1/p)\n",
        "    print(\"Using graduated nonconvexity in the loss\")\n",
        "  else:\n",
        "    data_loss = (jnp.abs(diff) ** 2).sum()\n",
        "    print(\"Using L2 loss\")\n",
        "\n",
        "  data_loss = data_loss / img_inds.shape[0] / spatial_inds.shape[-1]\n",
        "\n",
        "  loss = data_loss\n",
        "  loss += 0.1 * eikonal_loss / img_inds.shape[0]\n",
        "  loss += 1e-5 * (mask_area_loss * jnp.sin(omega_theta)).sum() * dtheta_dphi / img_inds.shape[0] / 4.0 / jnp.pi\n",
        "  #loss += 1e-1 * ((jnp.abs(params_envmap) ** 2) * ell.flatten()[:, None, None]).mean()\n",
        "  #loss += 0.01 * (entropy_loss * jnp.sin(omega_theta)).sum() * dtheta_dphi / img_inds.shape[0] / 4.0 / jnp.pi\n",
        "  #loss += 0.01 * length_loss / img_inds.shape[0]\n",
        "  return loss, (data_loss, eikonal_loss)\n",
        "\n",
        "def safe_exp(x):\n",
        "  return jnp.exp(jnp.minimum(x, 80.0))\n",
        "\n",
        "def tonemap_and_clip(x):\n",
        "  return np.clip(linear_to_srgb(x), 0.0, 1.0)\n",
        "\n",
        "def params_to_envmap(params_envmap):\n",
        "  #envmap = jax.nn.sigmoid(params_envmap)\n",
        "  #envmap = jax.nn.softplus(params_envmap)\n",
        "  if envmap_representation == 'SH':\n",
        "    #params_envmap = jnp.where(ell.flatten()[:, None, None] \u003c 5, params_envmap, 0.0)\n",
        "    envmap = jax.nn.softplus(jax.vmap(isht, in_axes=-1, out_axes=-1)(params_envmap))\n",
        "  else:\n",
        "    envmap = safe_exp(params_envmap)\n",
        "  return envmap\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def update_params(i, rng, state_envmap, state_sdf, state_materials, gt, spatial_inds, img_inds):\n",
        "  params_envmap = get_params_envmap(state_envmap)\n",
        "  params_sdf = get_params_sdf(state_sdf)\n",
        "  params_materials = get_params_materials(state_materials)\n",
        "\n",
        "  (loss, (data_loss, eikonal_loss)), g = jax.value_and_grad(get_loss, argnums=(0, 1, 2), has_aux=True)(params_envmap, params_sdf, params_materials,\n",
        "                                                                                                       gt, spatial_inds, img_inds, i, rng)\n",
        "\n",
        "  grad_envmap = jax.lax.pmean(g[0], axis_name='batch')\n",
        "  grad_sdf = jax.lax.pmean(g[1], axis_name='batch')\n",
        "  grad_materials = jax.lax.pmean(g[2], axis_name='batch')\n",
        "  eikonal_loss = jax.lax.pmean(eikonal_loss, axis_name='batch')\n",
        "  data_loss = jax.lax.pmean(data_loss, axis_name='batch')\n",
        "  loss = jax.lax.pmean(loss, axis_name='batch')\n",
        "\n",
        "  return update_envmap(i, grad_envmap, state_envmap), update_sdf(i, grad_sdf, state_sdf), update_materials(i, grad_materials, state_materials), loss, data_loss, eikonal_loss\n",
        "\n",
        "\n",
        "slow_optimization_mode = ('masks' not in gt_list and 'sdfs' not in gt_list) or 'materials' not in gt_list\n",
        "if slow_optimization_mode:\n",
        "  print(\"Slow optimization\")\n",
        "else:\n",
        "  print(\"Fast optimization\")\n",
        "\n",
        "num_iters = 200000 #50000 if slow_optimization_mode else 10000\n",
        "#num_iters = 50000\n",
        "straight_through_mode = 'soft'\n",
        "assert straight_through_mode in ['none', 'hard', 'soft']\n",
        "\n",
        "# TODO:\n",
        "# 1. It looks like using straight-through on the masks (with constant width 10) makes them be a little off,\n",
        "#    but improves the envmap (making it a little noisier because of the bad masks). Why?\n",
        "# 2. \n",
        "\n",
        "envmap_representation = 'direct'\n",
        "if envmap_representation == 'SH':\n",
        "  #params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_H, 3)) - 0.5) * 0.01\n",
        "  params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_H, 3, 2)) - 0.5) / (1 + ell.flatten()[:, None, None, None]) * 0.01\n",
        "  params_envmap = params_envmap[..., 0] + 1j * params_envmap[..., 1]\n",
        "  init_lr_envmap = 0.0003 if slow_optimization_mode else 0.001\n",
        "\n",
        "elif envmap_representation == 'direct':\n",
        "  params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_W, 3)) - 0.5) * 0.1 - 4.0 * int(slow_optimization_mode)\n",
        "  print(\"TODO: Get rid of this annoying -4.0\")\n",
        "\n",
        "  init_lr_envmap = 0.0003 if slow_optimization_mode else 0.03\n",
        "else:\n",
        "  raise ValueError('')\n",
        "  \n",
        "init_envmap, update_envmap, get_params_envmap = jax.experimental.optimizers.adam(init_lr_envmap)\n",
        "state_envmap = init_envmap(params_envmap)\n",
        "\n",
        "\n",
        "sdf_representation = 'pyramid'  # 'pyramid', 'mlp', 'pyramid'\n",
        "\n",
        "if sdf_representation == 'mlp':\n",
        "  L_encoding_sdf = 0 # 4\n",
        "  mlp_input_features = 3 + 6 * L_encoding_sdf\n",
        "\n",
        "  sdf_mlp = MLP([128]*4 + [1])\n",
        "\n",
        "  params_sdf = sdf_mlp.init(jax.random.PRNGKey(0),\n",
        "                            np.zeros([1, mlp_input_features]),\n",
        "                            np.zeros([1, 3]))\n",
        "  init_lr_sdf = 0.0001\n",
        "elif sdf_representation == 'pyramid':\n",
        "  pyramid_num_scales = 5\n",
        "  pyramid_resize_scale = 2\n",
        "  pyramid_mult = 2.0\n",
        "  global_std = 1.0 #0.1\n",
        "\n",
        "  rng = jax.random.PRNGKey(0)\n",
        "  params_sdf = get_pyramid_params(rng, N_cameras, envmap_H, envmap_W, pyramid_num_scales, pyramid_resize_scale, global_std)\n",
        "  init_lr_sdf = 0.003 #* 100\n",
        "  mask_width = 0.1\n",
        "\n",
        "elif sdf_representation == 'grid':\n",
        "  params_sdf = jax.random.normal(jax.random.PRNGKey(111), shape=(N_cameras, envmap_H, envmap_W))\n",
        "  init_lr_sdf = 0.01\n",
        "else:\n",
        "  raise ValueError('')\n",
        "\n",
        "#for _ in range(20):\n",
        "#  print(\"TODO: Create xmanager profile for this!\")\n",
        "#  print(\"TODO: Use xmanager to optimize: learning rates, biases, global std for params_sdf in pyramid mode, number of iterations (longer!), etc.\")\n",
        "#  print(\"TODO: Find out what happens if for masks we use the ground truth SDFs passed through a 0.1 sigmoid, instead of the GT masks. This is the best case scenario when using such a soft sigmoid!\")\n",
        "#  print(\"TODO: Replace initialization as soft ~0.5ish masks and dark envmap with good init. of envmap and ~1 masks (no occluders). Currently things just go there anyway...\")\n",
        "\n",
        "init_sdf, update_sdf, get_params_sdf = jax.experimental.optimizers.adam(init_lr_sdf)\n",
        "state_sdf = init_sdf(params_sdf)\n",
        "\n",
        "# Initialize material MLP\n",
        "L_encoding_materials = 0\n",
        "mlp_input_features = 3 + 6 * L_encoding_materials\n",
        "material_mlp = MLP([128]*4 + [3 + 6 * int(shading in ['phong', 'blinnphong'])])\n",
        "\n",
        "params_materials = material_mlp.init(jax.random.PRNGKey(0),\n",
        "                          np.zeros([1, mlp_input_features]))\n",
        "init_lr_materials = 0.0001\n",
        "init_materials, update_materials, get_params_materials = jax.experimental.optimizers.adam(init_lr_materials)\n",
        "state_materials = init_materials(params_materials)\n",
        "\n",
        "\n",
        "np_rng = np.random.default_rng(12345)\n",
        "jax_rng = jax.random.PRNGKey(3948)\n",
        "\n",
        "spatial_batch_size = 64\n",
        "for _ in range(10):\n",
        "  print(\"Try larger spatial batch size for optimizing materials!\")\n",
        "img_batch_size = 1024\n",
        "#img_batch_size = 64\n",
        "losses = []\n",
        "data_losses = []\n",
        "eikonal_losses = []\n",
        "envmap_psnrs = []\n",
        "tonemapped_envmap_psnrs = []\n",
        "mask_errors = []\n",
        "envmaps = []\n",
        "\n",
        "t = 0.0\n",
        "training_progress_bar = ProgressBar()\n",
        "training_progress_bar.Publish()\n",
        "\n",
        "replicated_state_envmap = flax.jax_utils.replicate(state_envmap)\n",
        "replicated_state_sdf = flax.jax_utils.replicate(state_sdf)\n",
        "replicated_state_materials = flax.jax_utils.replicate(state_materials)\n",
        "replicated_imgs_gt = flax.jax_utils.replicate(imgs_gt_noisy)\n",
        "\n",
        "\n",
        "for iteration in range(num_iters):\n",
        "  #t0 = time.time()\n",
        "  # Generate B1 image indices\n",
        "  img_inds = np_rng.choice(imgs_gt.shape[0], size=img_batch_size, replace=False) #* 0\n",
        "  # Now generate B2 pixel indices for each image. The total batch size is B1 * B2.\n",
        "\n",
        "  keys = jax.random.split(jax_rng, num=img_batch_size+1)\n",
        "  jax_rng, keys = keys[0], keys[1:]\n",
        "  spatial_inds = jax.vmap(jax.random.choice, in_axes=(0, None, None, None, 0))(keys, H*W, (spatial_batch_size,), False, alpha_gt[img_inds, :, 0])\n",
        "\n",
        "  assert jnp.all(alpha_gt[img_inds[:, None], spatial_inds[img_inds], :] \u003e 0.99)\n",
        "\n",
        "  replicated_state_envmap, replicated_state_sdf, replicated_state_materials, loss, data_loss, eikonal_loss = jax.pmap(update_params, in_axes=(None, None, 0, 0, 0, 0, 0, 0), axis_name='batch')(\n",
        "      iteration,\n",
        "      jax_rng,\n",
        "      replicated_state_envmap,\n",
        "      replicated_state_sdf,\n",
        "      replicated_state_materials,\n",
        "      replicated_imgs_gt,\n",
        "      spatial_inds.reshape(num_devices, -1, spatial_batch_size),\n",
        "      img_inds.reshape(num_devices, -1)\n",
        "      )\n",
        "\n",
        "  if iteration % 100 == 0 or iteration == num_iters - 1:\n",
        "    envmap = params_to_envmap(get_params_envmap(replicated_state_envmap)[0])\n",
        "    envmaps.append(envmap)\n",
        "    mse = (jnp.sin(omega_theta)[:, None] * (exposure * envmap_gt - envmap).reshape(-1, 3) ** 2).sum() * dtheta_dphi / 4.0 / jnp.pi / 3.0\n",
        "    envmap_psnrs.append(mse_to_psnr(mse))\n",
        "    mse = (jnp.sin(omega_theta)[:, None] * (tonemap_and_clip(exposure * envmap_gt) - tonemap_and_clip(envmap)).reshape(-1, 3) ** 2).sum() * dtheta_dphi / 4.0 / jnp.pi / 3.0\n",
        "    tonemapped_envmap_psnrs.append(mse_to_psnr(mse))\n",
        "\n",
        "  if iteration % 1000 == 0 or iteration == num_iters - 1:\n",
        "    sdf_params = get_params_sdf(flax.jax_utils.unreplicate(replicated_state_sdf))\n",
        "    if sdf_representation == 'mlp':\n",
        "      #sdfs = sdf_mlp.apply(sdf_params, posenc(omega_xyz)[None, :, :].repeat(N_cameras, 0), rays_o_vec[:, None, 0, 0, :].repeat(omega_xyz.shape[0], 1))\n",
        "\n",
        "      inp1 = posenc(omega_xyz, L_encoding_sdf)[None, :, :].repeat(N_cameras, 0)\n",
        "      inp2 = rays_o_vec[:, None, 0, 0, :].repeat(omega_xyz.shape[0], 1)\n",
        "\n",
        "      inp1 = inp1.reshape(num_devices, inp1.shape[0]//num_devices, *inp1.shape[1:])\n",
        "      inp2 = inp2.reshape(num_devices, inp2.shape[0]//num_devices, *inp2.shape[1:])\n",
        "      sdfs = jax.pmap(lambda x, y: sdf_mlp.apply(sdf_params, x, y))(inp1, inp2)\n",
        "    else:\n",
        "      sdfs = params_to_sdf(sdf_params, jnp.arange(N_cameras))\n",
        "\n",
        "    mask_error = (jnp.sin(omega_theta) * jnp.abs(masks_gt.reshape(N_cameras, -1) - (sdfs \u003e 0).reshape(N_cameras, -1))).sum() / N_cameras / 4.0 / jnp.pi\n",
        "    mask_errors.append(mask_error)\n",
        "\n",
        "  losses.append(loss[0])\n",
        "  data_losses.append(data_loss[0])\n",
        "  eikonal_losses.append(eikonal_loss[0])\n",
        "\n",
        "  if iteration in [500] or iteration % 1000 == 0 and iteration \u003e 0:\n",
        "    clear_output(wait=True)\n",
        "    training_progress_bar.Publish()\n",
        "    training_progress_bar.SetProgress(100.0 * (iteration + 1) / num_iters)\n",
        "\n",
        "    plt.figure()\n",
        "    plt.semilogy(data_losses)\n",
        "    plt.semilogy(eikonal_losses)\n",
        "\n",
        "    plt.figure(figsize=[16, 8])\n",
        "    plt.subplot(221)\n",
        "    plt.imshow(params_to_envmap(get_params_envmap(replicated_state_envmap)[0]))\n",
        "    plt.axis('off')\n",
        "    plt.subplot(222)\n",
        "    plt.imshow(exposure * envmap_gt)\n",
        "    plt.axis('off')\n",
        "    plt.subplot(223)\n",
        "    plt.imshow(linear_to_srgb(params_to_envmap(get_params_envmap(replicated_state_envmap)[0])))\n",
        "    plt.axis('off')\n",
        "    plt.subplot(224)\n",
        "    plt.imshow(linear_to_srgb(exposure * envmap_gt))\n",
        "    plt.axis('off')\n",
        "\n",
        "    plt.figure(figsize=[8, 8])\n",
        "    plt.subplot(121)\n",
        "    plt.plot(np.linspace(0, iteration, len(envmap_psnrs)), envmap_psnrs)\n",
        "    plt.subplot(122)\n",
        "    plt.plot(np.linspace(0, iteration, len(envmap_psnrs)), tonemapped_envmap_psnrs)\n",
        "\n",
        "    plt.figure()\n",
        "    plt.plot(np.linspace(0, iteration, len(mask_errors)), mask_errors)\n",
        "    # With known masks the loss gets to around 10^{-8}\n",
        "\n",
        "    plt.figure(figsize=[15, 12])\n",
        "    num_rows = 6\n",
        "    for row, i in enumerate([int(x) for x in np.linspace(0, N_cameras-1, num_rows)]):\n",
        "      sdf = sdfs.reshape(N_cameras, envmap_H, envmap_W)[i]\n",
        "      mask_gt = masks_gt[i]\n",
        "      diff = (jnp.float32((sdf \u003e 0) \u0026 (mask_gt \u003c 0.5)) - jnp.float32((sdf \u003c 0) \u0026 (mask_gt \u003e 0.5))) * 0.5 + 0.5\n",
        "      for col, img_to_plot in enumerate([sdf, sdf_to_mask(sdf, mask_width), sdf \u003e 0, mask_gt, diff]):\n",
        "        plt.subplot(num_rows, 5, row * 5 + col + 1)\n",
        "        if img_to_plot.shape[-1] != 3:\n",
        "          if col == 0:\n",
        "            plt.imshow(img_to_plot, cmap='gray')\n",
        "          else:\n",
        "            plt.imshow(img_to_plot, cmap='gray', vmin=0.0, vmax=1.0)\n",
        "        else:\n",
        "          plt.imshow(img_to_plot)\n",
        "        plt.axis('off')\n",
        "        plt.title(f'{i}')\n",
        "    plt.show()\n",
        "\n",
        "#sdfs.reshape(N_cameras, envmap_H, envmap_W).shape, masks_gt.shape\n",
        "\n",
        "#for _ in range(10):\n",
        "#  print(\"I don't understand how sigmoid with weight 0.1 like I'm using here is sufficient to get an envmap that's not a blurry mess???\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B43u_gSm_-5_"
      },
      "outputs": [],
      "source": [
        "mask_errors[-1]  # With 'none' straight through estimator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jvwh6RZYo8wb"
      },
      "outputs": [],
      "source": [
        "mask_errors[-1]  # With soft straight through estimator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RreSxFrpC3TD"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PAdnAjibzWvp"
      },
      "outputs": [],
      "source": [
        "params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_H, 3, 2)) - 0.5) / (1 + ell.flatten()[:, None, None, None]) * 0.2\n",
        "params_envmap = params_envmap[..., 0] + 1j * params_envmap[..., 1]\n",
        "\n",
        "envmap = safe_exp(jax.vmap(isht, in_axes=-1, out_axes=-1)(params_envmap) - 2.0)\n",
        "print(envmap.min(), envmap.max())\n",
        "plt.imshow(envmap)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fJkBEXfaRxdN"
      },
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WRf87kEvJkbs"
      },
      "outputs": [],
      "source": [
        "ind = N_cameras//2\n",
        "plt.imshow(params_to_sdf(get_params_sdf(flax.jax_utils.unreplicate(replicated_state_sdf)), jnp.arange(ind, ind+1))[0, :, :] \u003e -50.0)\n",
        "plt.figure()\n",
        "plt.imshow(masks_gt[ind])\n",
        "#plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GwrTP3J8qFj1"
      },
      "outputs": [],
      "source": [
        "plt.imshow(jnp.linalg.norm(params_to_envmap(get_params_envmap(replicated_state_envmap)[0]) / exposure / envmap_gt, axis=-1))\n",
        "plt.colorbar()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HAtNsLPacwGk"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gfHI4AgaB9yJ"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1W9oEJCat3Ga"
      },
      "outputs": [],
      "source": [
        "media.show_video(envmaps, height=400)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sxPT5ntFaqL8"
      },
      "outputs": [],
      "source": [
        "\n",
        "envmap = params_to_envmap(get_params_envmap(replicated_state_envmap)[0])\n",
        "mse = (jnp.sin(omega_theta)[:, None] * (f(envmap_gt) - f(envmap)).reshape(-1, 3) ** 2).sum() * dtheta_dphi / 4.0 / jnp.pi\n",
        "psnr = mse_to_psnr(mse)\n",
        "print(psnr)\n",
        "\n",
        "plt.imshow(linear_to_srgb(envmap))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4mRn5yLRnHoK"
      },
      "outputs": [],
      "source": [
        "def render(params_envmap, params_sdf, params_materials, img_ind):\n",
        "  img_inds = jnp.zeros((1,), dtype=jnp.int32) + img_ind\n",
        "  if 'envmap' in gt_list:\n",
        "    envmap = envmap_gt\n",
        "  else:\n",
        "    envmap = params_to_envmap(params_envmap)\n",
        "\n",
        "  if 'masks' in gt_list:\n",
        "    masks = masks_gt[img_inds]\n",
        "    eikonal_loss = 0.0\n",
        "    length_loss = 0.0\n",
        "  else:\n",
        "    sdf = params_to_sdf(params_sdf, img_inds)\n",
        "    sdf = sdf.reshape(1, envmap_H * envmap_W)\n",
        "\n",
        "    #mask_width = 10.0 #200.0 * (i / num_iters) + 10.0\n",
        "    #mask_width = 20.0 * (i / num_iters) + 10.0\n",
        "    #mask_width = 0.1\n",
        "    #mask_width = 10.0 #150.0 * (i / num_iters) + 6.0\n",
        "    curve = 'sigmoid'\n",
        "    if straight_through:\n",
        "      masks = sdf_to_mask(sdf, mask_width, curve)\n",
        "      masks = (masks + jax.lax.stop_gradient(jnp.float32(sdf \u003e 0.0) - masks)).reshape(img_inds.shape[0], envmap_H, envmap_W)\n",
        "    else:\n",
        "      masks = sdf_to_mask(sdf, mask_width, curve).reshape(img_inds.shape[0], envmap_H, envmap_W)\n",
        "\n",
        "  results = []\n",
        "  material_results = []\n",
        "  for i in range(H):\n",
        "    j0 = i * W\n",
        "    j1 = j0 + W\n",
        "    normals = normals_gt[img_inds, j0:j1, ...]\n",
        "    t_surface = t_surface_gt[img_inds, j0:j1, ...]\n",
        "    alpha = alpha_gt[img_inds, j0:j1, ...]\n",
        "    rays_d = rays_d_r[img_inds, j0:j1, ...]\n",
        "    rays_o = rays_o_r[img_inds, j0:j1, ...]\n",
        "\n",
        "    if 'materials' in gt_list:\n",
        "      materials = jax.tree_map(lambda x: x[img_ind, j0:j1, ...], materials_gt)\n",
        "    else:\n",
        "      pts = rays_o + rays_d * t_surface\n",
        "      materials = params_to_materials(params_materials, pts)\n",
        "\n",
        "\n",
        "    #print(materials['albedo'].shape, normals.shape, rays_d.shape, alpha.shape)\n",
        "    res = jax.vmap(render_gt_partial, in_axes=(None, 0, 0, 0, 0, 0, None, None))(\n",
        "        envmap,\n",
        "        masks,\n",
        "        materials,\n",
        "        normals,\n",
        "        rays_d,\n",
        "        alpha,\n",
        "        sphere_xyz,\n",
        "        sphere_rad,\n",
        "        )\n",
        "    results.append(res)\n",
        "    material_results.append(materials)\n",
        "\n",
        "  return results, material_results\n",
        "\n",
        "res, mat_res = render(get_params_envmap(flax.jax_utils.unreplicate(replicated_state_envmap)),\n",
        "             get_params_sdf(flax.jax_utils.unreplicate(replicated_state_sdf)),\n",
        "             get_params_materials(flax.jax_utils.unreplicate(replicated_state_materials)),\n",
        "             0)\n",
        "plt.imshow(jnp.concatenate(res, axis=0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KTN5Uc5Uv4g9"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RsD_y5Vxv4zE"
      },
      "outputs": [],
      "source": [
        "material_results = {}\n",
        "for k in mat_res[0].keys():\n",
        "  r = jnp.concatenate([m[k] for m in mat_res], axis=0)\n",
        "  material_results[k] = r\n",
        "  m, M = jnp.nanmin(r), jnp.nanmax(r)\n",
        "  r_normalized = (r - m) / (M - m)\n",
        "  plt.figure()\n",
        "  plt.imshow(r_normalized)\n",
        "  plt.title(f\"{k}: [{m:.3f}, {M:.3f}]\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xx5M_84sw27y"
      },
      "outputs": [],
      "source": [
        "plt.figure()\n",
        "plt.plot(material_results['albedo'][H//2, :, :])\n",
        "plt.figure()\n",
        "plt.plot(material_results['specular_albedo'][H//2, :, :])\n",
        "plt.figure()\n",
        "plt.plot(material_results['specular_exponent'][H//2, :, :])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "65HVO5acLj5p"
      },
      "outputs": [],
      "source": [
        "params_to_materials(get_params_materials(state_materials), jnp.zeros((10, 8, 3)))['albedo'].shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7bVKNRGYdjVk"
      },
      "outputs": [],
      "source": [
        "mlp_res = material_mlp.apply(get_params_materials(state_materials), posenc(jnp.zeros((10, 8, 3)), L_encoding_materials))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jeC0FH0GpAHw"
      },
      "outputs": [],
      "source": [
        "normals = normals_gt[0]\n",
        "normals = normals + jax.random.normal(jax.random.PRNGKey(0), shape=normals.shape) * 0.03\n",
        "normals = normals / jnp.linalg.norm(normals, axis=-1, keepdims=True)\n",
        "\n",
        "plt.imshow(normals.reshape(H, W, 3) * 0.5 + 0.5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "URJTcUOo9rdb"
      },
      "outputs": [],
      "source": [
        "num_devices = jax.local_device_count()\n",
        "\n",
        "#gt_list = ['materials']\n",
        "#gt_list = ['masks', 'materials']\n",
        "#gt_list = ['sdfs', 'materials']\n",
        "#gt_list = ['envmap', 'materials']\n",
        "#gt_list = ['envmap', 'masks']\n",
        "#gt_list = ['masks']\n",
        "gt_list = []\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  features: Sequence[int]\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x, y=None):\n",
        "    if y is not None:\n",
        "      x = jnp.concatenate([x, y], axis=-1)\n",
        "    for feat in self.features[:-1]:\n",
        "      x = nn.relu(nn.Dense(feat)(x))\n",
        "    x = nn.Dense(self.features[-1])(x)\n",
        "    return x\n",
        "\n",
        "def get_pyramid_params(rng, num_pyramids, height, width, pyramid_num_scales, pyramid_resize_scale, global_std, global_bias):\n",
        "  pyramid_params = []\n",
        "  for i in range(pyramid_num_scales):\n",
        "    gsh, gsw = [sz // pyramid_resize_scale ** i for sz in [height, width]]\n",
        "    key, rng = jax.random.split(rng)\n",
        "    features = jax.random.normal(key, (num_pyramids, gsh, gsw)) * global_std + global_bias\n",
        "    pyramid_params.append(features)\n",
        "  return pyramid_params\n",
        "\n",
        "\n",
        "def pyramids_to_imgs(pyramid_params, pyramid_mult, img_inds):\n",
        "  def pyramid_to_img(pyramid_params, pyramid_mult):\n",
        "    acc_val = pyramid_params[-1]\n",
        "    for i, curr_val in enumerate(pyramid_params[-2::-1], start=1):\n",
        "      # upsample\n",
        "      acc_val = jax.image.resize(acc_val * pyramid_mult, shape=curr_val.shape, method='linear')\n",
        "      \n",
        "      # accumulate\n",
        "      acc_val += curr_val\n",
        "    return acc_val\n",
        "\n",
        "  # Select all pyramid parameters at given indices\n",
        "  sub_pyramid_params = jax.tree_map(lambda t: t[img_inds], pyramid_params)\n",
        "  return jax.vmap(pyramid_to_img, in_axes=(0, None))(sub_pyramid_params, pyramid_mult)\n",
        "\n",
        "\n",
        "def grad_norm_spherical(dirs, grad):\n",
        "  \"\"\"\n",
        "  Compute gradient norm restricted to the sphere.\n",
        "  Assume dim 0 is elevation and 1 is azimuth.\n",
        "\n",
        "  dirs: (N, 3) array of directions on the sphere\n",
        "  grad: (N, 3) array of Cartesian gradients of points on dirs\n",
        "  \"\"\"\n",
        "\n",
        "  norm_spherical = grad - (dirs * grad).sum(axis=-1, keepdims=True) * dirs\n",
        "  return jnp.sqrt((norm_spherical ** 2).sum(-1) + 1e-5)\n",
        "\n",
        "\n",
        "def add_fake_quantization_noise(x, num_bits):\n",
        "  m = jnp.nanmin(x)\n",
        "  M = jnp.nanmax(x)\n",
        "  x = (x - m) / (M - m)\n",
        "  #return jnp.float32(jnp.uint8(x * 255.0)) / 255 * (M - m) + m\n",
        "  if num_bits == 8:\n",
        "    return np.float32(np.uint8(x * 255.0)) / 255.0 * (M - m) + m\n",
        "    \n",
        "  d = 0.5 ** num_bits  # This is the distance between quantization bins, and the bins' widths\n",
        "  x = 0.5 * d + (1.0 - d) * x\n",
        "  return (x + (np.random.rand(*x.shape) - 0.5) * d) * (M - m) + m\n",
        "\n",
        "rays_o_r = rays_o_vec.reshape(-1, H*W, 3)\n",
        "rays_d_r = rays_d_vec.reshape(-1, H*W, 3)\n",
        "\n",
        "num_bits = 1024\n",
        "imgs_gt_noisy = add_fake_quantization_noise(imgs_gt, num_bits)\n",
        "for _ in range(10):\n",
        "  print(f\"Using {num_bits} bits\")\n",
        "\n",
        "\n",
        "append_identity = True\n",
        "def posenc(x, L_encoding):\n",
        "  if L_encoding \u003c= 0:\n",
        "    return x\n",
        "  else:\n",
        "    scales = 2**jnp.arange(L_encoding)\n",
        "    #shape = x.shape[:-1] + (-1,)\n",
        "    #scaled_x = jnp.reshape((x[..., None, :] * scales[:, None]), shape)\n",
        "\n",
        "    #four_feat = jnp.sin(\n",
        "    #    jnp.concatenate([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1))\n",
        "    shape = x.shape[:-1] + (-1,)\n",
        "    scaled_x = x[..., None, :] * scales[:, None] # [..., L, D]\n",
        "\n",
        "    four_feat = jnp.sin(\n",
        "        jnp.stack([scaled_x, scaled_x + 0.5 * jnp.pi], axis=-1)) # [..., L, D, 2]\n",
        "\n",
        "    four_feat = jnp.reshape(four_feat / scales[:, None, None], shape)\n",
        "    print(\"Using Lipschitz posenc\")\n",
        "    if append_identity:\n",
        "      return jnp.concatenate([x] + [four_feat], axis=-1)\n",
        "    else:\n",
        "      return four_feat\n",
        "\n",
        "\n",
        "def params_to_sdf(params_sdf, img_inds):\n",
        "  if sdf_representation == 'mlp':\n",
        "    sdf, sdf_grad = jax.vmap(jax.value_and_grad(lambda x, y: sdf_mlp.apply(params_sdf, x, y)[0]))(\n",
        "        posenc(omega_xyz[None, :, :], L_encoding_sdf).repeat(img_inds.shape[0], 0).reshape(-1, mlp_input_features),\n",
        "        rays_o_vec[img_inds, 0, 0, :][:, None, :].repeat(envmap_H * envmap_W, 1).reshape(-1, 3)\n",
        "        )\n",
        "  elif sdf_representation == 'grid':\n",
        "    sdf = params_sdf[img_inds[:, None], ...]\n",
        "  else:\n",
        "    sdf = pyramids_to_imgs(params_sdf, pyramid_mult, img_inds)\n",
        "  return sdf\n",
        "\n",
        "\n",
        "def sdf_to_mask(x, width, curve='sigmoid'):\n",
        "  if curve == 'sigmoid':\n",
        "    return jax.nn.sigmoid(x * width)\n",
        "  elif curve == 'laplace_cdf':\n",
        "    return 0.5 + 0.5 * jnp.sign(x) * (1.0 - jnp.exp(-jnp.abs(x) * width))\n",
        "  else:\n",
        "    raise NotImplementedError('Only sigmoid and laplace_cdf for now.')\n",
        "  \n",
        "def params_to_materials(params_materials, pts):\n",
        "  mlp_res = material_mlp.apply(params_materials, posenc(pts, L_encoding_materials))\n",
        "  materials = {}\n",
        "  materials['albedo'] = jax.nn.sigmoid(mlp_res[..., 0:3])\n",
        "  if shading in ['phong', 'blinnphong']:\n",
        "    #materials['specular_albedo'] = 20.0 * jax.nn.sigmoid(mlp_res[..., 3:6])\n",
        "    if is_dielectric:\n",
        "      materials['specular_albedo'] = jax.nn.softplus(mlp_res[..., 3:4])\n",
        "      materials['specular_exponent'] = jax.nn.softplus(mlp_res[..., 4:5])\n",
        "    else:\n",
        "      materials['specular_albedo'] = jax.nn.softplus(mlp_res[..., 3:6])\n",
        "      materials['specular_exponent'] = jax.nn.softplus(mlp_res[..., 6:7])\n",
        "\n",
        "  return materials\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def get_loss(params_envmap, params_sdf, params_materials, gt, spatial_inds, img_inds, i, rng):\n",
        "  rng, key = jax.random.split(rng)\n",
        "\n",
        "  if 'envmap' in gt_list:\n",
        "    envmap = envmap_gt * exposure\n",
        "  else:\n",
        "    envmap = params_to_envmap(params_envmap)\n",
        "\n",
        "  normals = normals_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  t_surface = t_surface_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  alpha = alpha_gt[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  rays_d = rays_d_r[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "  rays_o = rays_o_r[img_inds[:, None], spatial_inds[img_inds], ...]\n",
        "\n",
        "  if False:\n",
        "    print(\"Adding normal noise\")\n",
        "    key, key2 = jax.random.split(key)\n",
        "    normals = normals + jax.random.normal(key2, shape=normals.shape) * 0.03\n",
        "    normals = normals / jnp.linalg.norm(normals, axis=-1, keepdims=True)\n",
        "\n",
        "  if 'materials' in gt_list:\n",
        "    materials = jax.tree_map(lambda x: x[img_inds[:, None], spatial_inds[img_inds], ...], materials_gt)\n",
        "  else:\n",
        "    pts = rays_o + rays_d * t_surface\n",
        "    materials = params_to_materials(params_materials, pts)\n",
        "\n",
        "  sdf = params_to_sdf(params_sdf, img_inds)\n",
        "  sdf = sdf.reshape(img_inds.shape[0], envmap_H * envmap_W)\n",
        "  #mask_width = 200.0 * (i / num_iters) + 10.0\n",
        "  #masks = sdf_to_mask(sdf.reshape(-1, envmap_H, envmap_W), mask_width, 'sigmoid')\n",
        "\n",
        "  sdf_curve = 'sigmoid'\n",
        "\n",
        "  if 'masks' in gt_list or 'sdfs' in gt_list:\n",
        "    if 'masks' in gt_list:\n",
        "      masks = masks_gt[img_inds]\n",
        "    else:\n",
        "      masks = sdf_to_mask(sdfs_gt[img_inds], mask_width, sdf_curve)\n",
        "    eikonal_loss = 0.0\n",
        "    length_loss = 0.0\n",
        "    mask_area_loss = 0.0\n",
        "  else:\n",
        "    #mask_width = 10.0 #200.0 * (i / num_iters) + 10.0\n",
        "    #mask_width = 20.0 * (i / num_iters) + 10.0\n",
        "    #mask_width = 0.1\n",
        "    #mask_width = 10.0 #150.0 * (i / num_iters) + 6.0\n",
        "    #mask_width = 0.1 * jnp.exp(jnp.log(100) * i / num_iters)\n",
        "    if straight_through_mode == 'hard':\n",
        "      masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "      masks = (masks + jax.lax.stop_gradient(jnp.float32(sdf \u003e 0.0) - masks))\n",
        "    elif straight_through_mode == 'soft':\n",
        "      soft_masks = sdf_to_mask(sdf, 0.1, sdf_curve)\n",
        "      hard_masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "      # Define masks with value of `hard_masks` but gradients of `soft_masks`\n",
        "      masks = (soft_masks + jax.lax.stop_gradient(hard_masks - soft_masks))\n",
        "    elif straight_through_mode == 'none':\n",
        "      masks = sdf_to_mask(sdf, mask_width, sdf_curve)\n",
        "\n",
        "\n",
        "\n",
        "    if sdf_representation == 'mlp':\n",
        "      # TODO: Only compute grad w.r.t. x, not w.r.t. other posenc components\n",
        "      sdf_grad = sdf_grad[..., :3].reshape(img_inds.shape[0], envmap_H * envmap_W, 3)\n",
        "      sdf_grad_norm = grad_norm_spherical(omega_xyz, sdf_grad.reshape(img_inds.shape[0], envmap_H * envmap_W, 3))\n",
        "      eikonal_loss = (jnp.sin(omega_theta) * (sdf_grad_norm - 1) ** 2).sum() * dtheta_dphi\n",
        "    else:\n",
        "      eikonal_loss = 0.0\n",
        "\n",
        "    # Compute entropy assuming mask = sigmoid(sdf)\n",
        "    entropy_loss = jax.nn.softplus(-sdf) + sdf * (1.0 - jax.nn.sigmoid(sdf))\n",
        "    mask_area_loss = 1.0 - masks  # Try to make the occluder as small as possible\n",
        "    #delta_sdf = jax.vmap(jax.grad(lambda x: sdf_to_mask(x, mask_width, sdf_curve)))(sdf.flatten()).reshape(img_inds.shape[0], envmap_H * envmap_W)\n",
        "\n",
        "    #length_loss = (jnp.sin(omega_theta) * sdf_grad_norm * delta_sdf).sum() * dtheta_dphi\n",
        "\n",
        "  res = jax.vmap(render_gt_partial, in_axes=(None, 0, 0, 0, 0, 0, None, None))(\n",
        "      envmap,\n",
        "      masks.reshape(img_inds.shape[0], envmap_H, envmap_W),\n",
        "      materials,\n",
        "      normals,\n",
        "      rays_d,\n",
        "      alpha,\n",
        "      sphere_xyz,\n",
        "      sphere_rad,\n",
        "      )\n",
        "\n",
        "  #diff = gt[img_inds[:, None], spatial_inds[img_inds], :] - res\n",
        "  diff = linear_to_srgb(gt[img_inds[:, None], spatial_inds[img_inds], :]) - linear_to_srgb(res)\n",
        "  #loss_per_element = (diff ** 2).sum(-1)\n",
        "  #loss_per_element = jnp.abs(diff).sum(-1)\n",
        "  if False:\n",
        "    p = 2 - 1.5 * i / num_iters\n",
        "    data_loss = jnp.power((jnp.abs(diff + 1e-10) ** p).sum(), 1/p)\n",
        "    print(\"Using graduated nonconvexity in the loss\")\n",
        "  else:\n",
        "    data_loss = (jnp.abs(diff) ** 2).sum()\n",
        "    print(\"Using L2 loss\")\n",
        "\n",
        "  data_loss = data_loss / img_inds.shape[0] / spatial_inds.shape[-1]\n",
        "\n",
        "  loss = data_loss\n",
        "  loss += 0.1 * eikonal_loss / img_inds.shape[0]\n",
        "  loss += 1e-5 * (mask_area_loss * jnp.sin(omega_theta)).sum() * dtheta_dphi / img_inds.shape[0] / 4.0 / jnp.pi\n",
        "  #loss += 1e-1 * ((jnp.abs(params_envmap) ** 2) * ell.flatten()[:, None, None]).mean()\n",
        "  #loss += 0.01 * (entropy_loss * jnp.sin(omega_theta)).sum() * dtheta_dphi / img_inds.shape[0] / 4.0 / jnp.pi\n",
        "  #loss += 0.01 * length_loss / img_inds.shape[0]\n",
        "\n",
        "  # Environment map TV loss\n",
        "  loss += 1e-7 * (((envmap[:, 1:] - envmap[:, :-1]) ** 2).sum() + ((envmap[1:, :] - envmap[:-1, :]) ** 2).sum())  * dtheta_dphi / 4.0 / jnp.pi\n",
        "\n",
        "  return loss, (data_loss, eikonal_loss)\n",
        "\n",
        "def safe_exp(x):\n",
        "  return jnp.exp(jnp.minimum(x, 80.0))\n",
        "\n",
        "def tonemap_and_clip(x):\n",
        "  return np.clip(linear_to_srgb(x), 0.0, 1.0)\n",
        "\n",
        "def params_to_envmap(params_envmap):\n",
        "  #envmap = jax.nn.sigmoid(params_envmap)\n",
        "  #envmap = jax.nn.softplus(params_envmap)\n",
        "  if envmap_representation == 'SH':\n",
        "    #params_envmap = jnp.where(ell.flatten()[:, None, None] \u003c 5, params_envmap, 0.0)\n",
        "    envmap = jax.nn.softplus(jax.vmap(isht, in_axes=-1, out_axes=-1)(params_envmap))\n",
        "  else:\n",
        "    envmap = safe_exp(params_envmap)\n",
        "  return envmap\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def update_params(i, rng, state_envmap, state_sdf, state_materials, gt, spatial_inds, img_inds):\n",
        "  params_envmap = get_params_envmap(state_envmap)\n",
        "  params_sdf = get_params_sdf(state_sdf)\n",
        "  params_materials = get_params_materials(state_materials)\n",
        "\n",
        "  (loss, (data_loss, eikonal_loss)), g = jax.value_and_grad(get_loss, argnums=(0, 1, 2), has_aux=True)(params_envmap, params_sdf, params_materials,\n",
        "                                                                                                       gt, spatial_inds, img_inds, i, rng)\n",
        "\n",
        "  grad_envmap = jax.lax.pmean(g[0], axis_name='batch')\n",
        "  grad_sdf = jax.lax.pmean(g[1], axis_name='batch')\n",
        "  grad_materials = jax.lax.pmean(g[2], axis_name='batch')\n",
        "  eikonal_loss = jax.lax.pmean(eikonal_loss, axis_name='batch')\n",
        "  data_loss = jax.lax.pmean(data_loss, axis_name='batch')\n",
        "  loss = jax.lax.pmean(loss, axis_name='batch')\n",
        "\n",
        "  return update_envmap(i, grad_envmap, state_envmap), update_sdf(i, grad_sdf, state_sdf), update_materials(i, grad_materials, state_materials), loss, data_loss, eikonal_loss\n",
        "\n",
        "\n",
        "#slow_optimization_mode = ('masks' not in gt_list and 'sdfs' not in gt_list) or 'materials' not in gt_list\n",
        "#if slow_optimization_mode:\n",
        "#  print(\"Slow optimization\")\n",
        "#else:\n",
        "#  print(\"Fast optimization\")\n",
        "\n",
        "num_iters = 150000 #50000 if slow_optimization_mode else 10000\n",
        "#num_iters = 50000\n",
        "straight_through_mode = 'soft'\n",
        "assert straight_through_mode in ['none', 'hard', 'soft']\n",
        "\n",
        "# TODO:\n",
        "# 1. It looks like using straight-through on the masks (with constant width 10) makes them be a little off,\n",
        "#    but improves the envmap (making it a little noisier because of the bad masks). Why?\n",
        "# 2. \n",
        "\n",
        "envmap_representation = 'direct'\n",
        "if envmap_representation == 'SH':\n",
        "  #params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_H, 3)) - 0.5) * 0.01\n",
        "  params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_H, 3, 2)) - 0.5) / (1 + ell.flatten()[:, None, None, None]) * 0.01\n",
        "  params_envmap = params_envmap[..., 0] + 1j * params_envmap[..., 1]\n",
        "  init_lr_envmap = 0.0003 #if slow_optimization_mode else 0.001\n",
        "\n",
        "elif envmap_representation == 'direct':\n",
        "  params_envmap = (jax.random.uniform(jax.random.PRNGKey(0), shape=(envmap_H, envmap_W, 3)) - 0.5) * 0.1\n",
        "  #print(\"TODO: Get rid of this annoying -4.0\")\n",
        "\n",
        "  init_lr_envmap = 0.03 #if slow_optimization_mode else 0.03\n",
        "else:\n",
        "  raise ValueError('')\n",
        "  \n",
        "init_envmap, update_envmap, get_params_envmap = jax.experimental.optimizers.adam(init_lr_envmap)\n",
        "state_envmap = init_envmap(params_envmap)\n",
        "\n",
        "\n",
        "sdf_representation = 'pyramid'  # 'pyramid', 'mlp', 'pyramid'\n",
        "\n",
        "if sdf_representation == 'mlp':\n",
        "  L_encoding_sdf = 0 # 4\n",
        "  mlp_input_features = 3 + 6 * L_encoding_sdf\n",
        "\n",
        "  sdf_mlp = MLP([128]*4 + [1])\n",
        "\n",
        "  params_sdf = sdf_mlp.init(jax.random.PRNGKey(0),\n",
        "                            np.zeros([1, mlp_input_features]),\n",
        "                            np.zeros([1, 3]))\n",
        "  init_lr_sdf = 0.0001\n",
        "elif sdf_representation == 'pyramid':\n",
        "  pyramid_num_scales = 5\n",
        "  pyramid_resize_scale = 2\n",
        "  pyramid_mult = 2.0\n",
        "  global_std = 1.0 #0.1\n",
        "  global_bias = 4.0\n",
        "\n",
        "  rng = jax.random.PRNGKey(0)\n",
        "  params_sdf = get_pyramid_params(rng, N_cameras, envmap_H, envmap_W, pyramid_num_scales, pyramid_resize_scale, global_std, global_bias)\n",
        "  init_lr_sdf = 0.1 #* 100\n",
        "  mask_width = 0.1\n",
        "\n",
        "elif sdf_representation == 'grid':\n",
        "  params_sdf = jax.random.normal(jax.random.PRNGKey(111), shape=(N_cameras, envmap_H, envmap_W))\n",
        "  init_lr_sdf = 0.01\n",
        "else:\n",
        "  raise ValueError('')\n",
        "\n",
        "#for _ in range(20):\n",
        "#  print(\"TODO: Use xmanager to optimize: learning rates, biases, global std for params_sdf in pyramid mode, number of iterations (longer!), etc.\")\n",
        "#  print(\"TODO: Find out what happens if for masks we use the ground truth SDFs passed through a 0.1 sigmoid, instead of the GT masks. This is the best case scenario when using such a soft sigmoid!\")\n",
        "#  print(\"TODO: Replace initialization as soft ~0.5ish masks and dark envmap with good init. of envmap and ~1 masks (no occluders). Currently things just go there anyway...\")\n",
        "\n",
        "init_sdf, update_sdf, get_params_sdf = jax.experimental.optimizers.adam(init_lr_sdf)\n",
        "state_sdf = init_sdf(params_sdf)\n",
        "\n",
        "# Initialize material MLP\n",
        "L_encoding_materials = 4\n",
        "mlp_input_features = 3 + 6 * L_encoding_materials\n",
        "\n",
        "is_dielectric = True\n",
        "num_components = 3  # 3 for diffuse\n",
        "if shading in ['phong', 'blinnphong']:\n",
        "  if is_dielectric:\n",
        "    num_components += 2  # 1 for specular albedo, 1 for exponent\n",
        "  else:\n",
        "    num_components += 4  # 3 for specular albedo, 1 for exponent\n",
        "material_mlp = MLP([128]*4 + [num_components])\n",
        "\n",
        "params_materials = material_mlp.init(jax.random.PRNGKey(0),\n",
        "                          np.zeros([1, mlp_input_features]))\n",
        "init_lr_materials = 0.003\n",
        "init_materials, update_materials, get_params_materials = jax.experimental.optimizers.adam(init_lr_materials)\n",
        "state_materials = init_materials(params_materials)\n",
        "\n",
        "\n",
        "np_rng = np.random.default_rng(12345)\n",
        "jax_rng = jax.random.PRNGKey(3948)\n",
        "\n",
        "spatial_batch_size = 64\n",
        "\n",
        "img_batch_size = 512 #1024\n",
        "#img_batch_size = 64\n",
        "losses = []\n",
        "data_losses = []\n",
        "eikonal_losses = []\n",
        "envmap_psnrs = []\n",
        "tonemapped_envmap_psnrs = []\n",
        "mask_errors = []\n",
        "envmaps = []\n",
        "\n",
        "t = 0.0\n",
        "training_progress_bar = ProgressBar()\n",
        "training_progress_bar.Publish()\n",
        "\n",
        "replicated_state_envmap = flax.jax_utils.replicate(state_envmap)\n",
        "replicated_state_sdf = flax.jax_utils.replicate(state_sdf)\n",
        "replicated_state_materials = flax.jax_utils.replicate(state_materials)\n",
        "replicated_imgs_gt = flax.jax_utils.replicate(imgs_gt_noisy)\n",
        "\n",
        "\n",
        "for iteration in range(num_iters):\n",
        "  #t0 = time.time()\n",
        "  # Generate B1 image indices\n",
        "  img_inds = np_rng.choice(imgs_gt.shape[0], size=img_batch_size, replace=False) #* 0\n",
        "  # Now generate B2 pixel indices for each image. The total batch size is B1 * B2.\n",
        "\n",
        "  keys = jax.random.split(jax_rng, num=img_batch_size+1)\n",
        "  jax_rng, keys = keys[0], keys[1:]\n",
        "  spatial_inds = jax.vmap(jax.random.choice, in_axes=(0, None, None, None, 0))(keys, H*W, (spatial_batch_size,), False, alpha_gt[img_inds, :, 0])\n",
        "\n",
        "  assert jnp.all(alpha_gt[img_inds[:, None], spatial_inds[img_inds], :] \u003e 0.99)\n",
        "\n",
        "  replicated_state_envmap, replicated_state_sdf, replicated_state_materials, loss, data_loss, eikonal_loss = jax.pmap(update_params, in_axes=(None, None, 0, 0, 0, 0, 0, 0), axis_name='batch')(\n",
        "      iteration,\n",
        "      jax_rng,\n",
        "      replicated_state_envmap,\n",
        "      replicated_state_sdf,\n",
        "      replicated_state_materials,\n",
        "      replicated_imgs_gt,\n",
        "      spatial_inds.reshape(num_devices, -1, spatial_batch_size),\n",
        "      img_inds.reshape(num_devices, -1)\n",
        "      )\n",
        "\n",
        "  if iteration % 100 == 0 or iteration == num_iters - 1:\n",
        "    envmap = params_to_envmap(get_params_envmap(replicated_state_envmap)[0])\n",
        "    envmaps.append(envmap)\n",
        "    mse = (jnp.sin(omega_theta)[:, None] * (exposure * envmap_gt - envmap).reshape(-1, 3) ** 2).sum() * dtheta_dphi / 4.0 / jnp.pi / 3.0\n",
        "    envmap_psnrs.append(mse_to_psnr(mse))\n",
        "    mse = (jnp.sin(omega_theta)[:, None] * (tonemap_and_clip(exposure * envmap_gt) - tonemap_and_clip(envmap)).reshape(-1, 3) ** 2).sum() * dtheta_dphi / 4.0 / jnp.pi / 3.0\n",
        "    tonemapped_envmap_psnrs.append(mse_to_psnr(mse))\n",
        "\n",
        "  if iteration % 1000 == 0 or iteration == num_iters - 1:\n",
        "    sdf_params = get_params_sdf(flax.jax_utils.unreplicate(replicated_state_sdf))\n",
        "    if sdf_representation == 'mlp':\n",
        "      #sdfs = sdf_mlp.apply(sdf_params, posenc(omega_xyz)[None, :, :].repeat(N_cameras, 0), rays_o_vec[:, None, 0, 0, :].repeat(omega_xyz.shape[0], 1))\n",
        "\n",
        "      inp1 = posenc(omega_xyz, L_encoding_sdf)[None, :, :].repeat(N_cameras, 0)\n",
        "      inp2 = rays_o_vec[:, None, 0, 0, :].repeat(omega_xyz.shape[0], 1)\n",
        "\n",
        "      inp1 = inp1.reshape(num_devices, inp1.shape[0]//num_devices, *inp1.shape[1:])\n",
        "      inp2 = inp2.reshape(num_devices, inp2.shape[0]//num_devices, *inp2.shape[1:])\n",
        "      sdfs = jax.pmap(lambda x, y: sdf_mlp.apply(sdf_params, x, y))(inp1, inp2)\n",
        "    else:\n",
        "      sdfs = params_to_sdf(sdf_params, jnp.arange(N_cameras))\n",
        "\n",
        "    mask_error = (jnp.sin(omega_theta) * jnp.abs(masks_gt.reshape(N_cameras, -1) - (sdfs \u003e 0).reshape(N_cameras, -1))).sum() / N_cameras / 4.0 / jnp.pi\n",
        "    mask_errors.append(mask_error)\n",
        "\n",
        "  losses.append(loss[0])\n",
        "  data_losses.append(data_loss[0])\n",
        "  eikonal_losses.append(eikonal_loss[0])\n",
        "\n",
        "  if iteration in [500] or iteration % 1000 == 0 and iteration \u003e 0:\n",
        "    clear_output(wait=True)\n",
        "    training_progress_bar.Publish()\n",
        "    training_progress_bar.SetProgress(100.0 * (iteration + 1) / num_iters)\n",
        "\n",
        "    plt.figure()\n",
        "    plt.semilogy(data_losses)\n",
        "    plt.semilogy(eikonal_losses)\n",
        "\n",
        "    # Plot envmaps\n",
        "    plt.figure(figsize=[16, 8])\n",
        "    plt.subplot(221)\n",
        "    plt.imshow(params_to_envmap(get_params_envmap(flax.jax_utils.unreplicate(replicated_state_envmap))))\n",
        "    plt.axis('off')\n",
        "    plt.subplot(222)\n",
        "    plt.imshow(exposure * envmap_gt)\n",
        "    plt.axis('off')\n",
        "    plt.subplot(223)\n",
        "    plt.imshow(linear_to_srgb(params_to_envmap(get_params_envmap(flax.jax_utils.unreplicate(replicated_state_envmap)))))\n",
        "    plt.axis('off')\n",
        "    plt.subplot(224)\n",
        "    plt.imshow(linear_to_srgb(exposure * envmap_gt))\n",
        "    plt.axis('off')\n",
        "\n",
        "    # Plot materials\n",
        "    ind = 10\n",
        "    pts = rays_o_r[ind] + t_surface_gt[ind] * rays_d_r[ind]\n",
        "    materials = params_to_materials(get_params_materials(flax.jax_utils.unreplicate(replicated_state_materials)), pts)\n",
        "    for k in materials.keys():\n",
        "      plt.figure()\n",
        "      plt.subplot(121)\n",
        "      plt.imshow(materials[k].reshape(H, W, -1))\n",
        "      plt.axis('off')\n",
        "      plt.subplot(122)\n",
        "      plt.imshow(materials_gt[k][ind].reshape(H, W, -1))\n",
        "      plt.axis('off')\n",
        "    \n",
        "    # Plot PSNRs\n",
        "    plt.figure(figsize=[8, 8])\n",
        "    plt.subplot(121)\n",
        "    plt.plot(np.linspace(0, iteration, len(envmap_psnrs)), envmap_psnrs)\n",
        "    plt.subplot(122)\n",
        "    plt.plot(np.linspace(0, iteration, len(envmap_psnrs)), tonemapped_envmap_psnrs)\n",
        "\n",
        "    # Plot mask errors and masks\n",
        "    plt.figure()\n",
        "    plt.plot(np.linspace(0, iteration, len(mask_errors)), mask_errors)\n",
        "\n",
        "    plt.figure(figsize=[15, 12])\n",
        "    num_rows = 6\n",
        "    for row, i in enumerate([int(x) for x in np.linspace(0, N_cameras-1, num_rows)]):\n",
        "      sdf = sdfs.reshape(N_cameras, envmap_H, envmap_W)[i]\n",
        "      mask_gt = masks_gt[i]\n",
        "      diff = (jnp.float32((sdf \u003e 0) \u0026 (mask_gt \u003c 0.5)) - jnp.float32((sdf \u003c 0) \u0026 (mask_gt \u003e 0.5))) * 0.5 + 0.5\n",
        "      for col, img_to_plot in enumerate([sdf, sdf_to_mask(sdf, mask_width), sdf \u003e 0, mask_gt, diff]):\n",
        "        plt.subplot(num_rows, 5, row * 5 + col + 1)\n",
        "        if img_to_plot.shape[-1] != 3:\n",
        "          if col == 0:\n",
        "            plt.imshow(img_to_plot, cmap='gray')\n",
        "          else:\n",
        "            plt.imshow(img_to_plot, cmap='gray', vmin=0.0, vmax=1.0)\n",
        "        else:\n",
        "          plt.imshow(img_to_plot)\n",
        "        plt.axis('off')\n",
        "        plt.title(f'{i}')\n",
        "    plt.show()\n",
        "\n",
        "#sdfs.reshape(N_cameras, envmap_H, envmap_W).shape, masks_gt.shape\n",
        "\n",
        "#for _ in range(10):\n",
        "#  print(\"I don't understand how sigmoid with weight 0.1 like I'm using here is sufficient to get an envmap that's not a blurry mess???\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QIRuXUkTjhy_"
      },
      "outputs": [],
      "source": [
        "Blinnphong!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CZYwnkqbkOg2"
      },
      "source": [
        "We want to find:\n",
        "$$s^* = \\text{argmin}||sI-I'||^2$$\n",
        "\n",
        "We have that the derivative of the objective w.r.t. $s$ is:\n",
        "$$\\frac{d}{ds}||sI-I'||^2 = \\frac{d}{ds}\\sum_{j} (sI_j - I'_j)^2 = \\sum_j (sI_j-I_j')I_j \\overset{!}{=}0$$\n",
        "\n",
        "and therefore:\n",
        "$$s^*=\\frac{\\sum_j I_j I'_j}{\\sum_j I_j^2}$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "88roshT9jjrA"
      },
      "outputs": [],
      "source": [
        "#jnp.mean(envmap_gt[:, :, 0:1] / envmap[:, :, 0:1])\n",
        "s_star = jnp.mean(envmap[:, :, 0:1]*envmap_gt[:, :, 0:1]) / jnp.mean(envmap[:, :, 0:1]**2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xCfPTTWzi1lr"
      },
      "outputs": [],
      "source": [
        "s = np.logspace(-1.0, 1.0, 1000)\n",
        "es = ((s[None, None, :] * envmap[:, :, 0:1] - envmap_gt[:, :, 0:1]) ** 2).mean(0).mean(0)\n",
        "print(s[es.argmin()])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p8DpVIaKlsaJ"
      },
      "outputs": [],
      "source": [
        "exposure / s_star"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ozTcdfkFlDT7"
      },
      "outputs": [],
      "source": [
        "#mse = ((s_star * envmap[:, :, 0:1] - envmap_gt[:, :, 0:1]) ** 2).mean()\n",
        "mse = (jnp.sin(omega_theta)[:, None] * (tonemap_and_clip(exposure * envmap_gt) - tonemap_and_clip(envmap)).reshape(-1, 3) ** 2).sum() * dtheta_dphi / 4.0 / jnp.pi\n",
        "\n",
        "psnr = -10. / jnp.log(10.) * jnp.log(mse)\n",
        "print(mse)\n",
        "print(psnr, 'dB')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ao8ZEq7fUzz1"
      },
      "outputs": [],
      "source": [
        "ind = 10\n",
        "pts = rays_o_r[ind] + t_surface_gt[ind] * rays_d_r[ind]\n",
        "materials = params_to_materials(get_params_materials(flax.jax_utils.unreplicate(replicated_state_materials)), pts)\n",
        "for k in materials.keys():\n",
        "  print(k)\n",
        "  plt.figure()\n",
        "  plt.subplot(121)\n",
        "  plt.imshow(materials[k].reshape(H, W, -1))\n",
        "  plt.axis('off')\n",
        "  plt.subplot(122)\n",
        "  plt.imshow(materials_gt[k][ind].reshape(H, W, -1))\n",
        "  plt.axis('off')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eE9vEpyDU9HV"
      },
      "outputs": [],
      "source": [
        "plt.imshow(materials['specular_albedo'].reshape(H, W, 1))\n",
        "print(jnp.nanmean(materials['specular_albedo']))\n",
        "plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SB9QWEq6nuVC"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "opRqy-B-18mx"
      },
      "outputs": [],
      "source": [
        "materials_gt['albedo'].shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7b_CncOs06lG"
      },
      "outputs": [],
      "source": [
        "plt.imshow(material_mlp.apply(get_params_materials(flax.jax_utils.unreplicate(replicated_state_materials)), pts).reshape(H, W, 3))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TY8RWGC0tf1T"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OmmdCSHhbpYV"
      },
      "outputs": [],
      "source": [
        "tonemapped_envmap_psnrs[-1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7KPW7T1OtxoR"
      },
      "outputs": [],
      "source": [
        "  import pickle\n",
        "\n",
        "trained_params = jax.experimental.optimizers.unpack_optimizer_state(replicated_state_sdf)\n",
        "pickle.dump(trained_params, open(\"replicated_state_sdf.pkl\", \"wb\"))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "02OOcsStv-xZ"
      },
      "outputs": [],
      "source": [
        "best_params = pickle.load(open(\"replicated_state_sdf.pkl\", \"rb\"))\n",
        "replicated_state_sdf = jax.experimental.optimizers.pack_optimizer_state(best_params)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t30B4ca1wUj_"
      },
      "outputs": [],
      "source": [
        "#plt.imshow(get_params_sdf(replicated_state_sdf)[0][0][9])\n",
        "plt.imshow(pyramids_to_imgs(get_params_sdf(replicated_state_sdf), pyramid_mult, jnp.arange(1))[0][0] \u003e 0)\n",
        "plt.figure()\n",
        "plt.imshow(masks_gt[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rqSm1QvwhGHn"
      },
      "outputs": [],
      "source": [
        "\"\""
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//googlex/gcam/buff/mipnerf360:notebook_minus_nerf",
        "kind": "private"
      },
      "name": "spherical_peter_problem_materials.ipynb",
      "private_outputs": true,
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
